Master the iFLYTEK Prohibited Words Classification Challenge: Baselines & BERT

Welcome from the author, who shares a competition related to content safety and provides several simple baselines for participants.

Competition Name

iFLYTEK AI Developer Competition – Prohibited Words Classification Challenge.

Competition Type

Natural Language Processing.

Task

Accurately identify the category of prohibited words in text.

Background

With the rapid growth of online content, platforms face increasing challenges in moderating prohibited information. Prohibited words (e.g., political sensitivity, sexual content, violence, superstition) appear in dynamic forms such as homophones, abbreviations, symbol insertions, and multilingual mixes, continuously evading traditional keyword‑based rules and threatening ecosystem safety and user experience. Existing keyword‑matching or simple rule‑based methods struggle with semantic ambiguity and adversarial interference, requiring AI techniques to improve detection intelligence.

Data Description

The competition provides a large, anonymized dataset collected from various social platforms, containing the main text and its corresponding prohibited‑word category for participants to use.

Evaluation Metric

Submissions are evaluated using macro F1‑score.

Baseline 1 – Logistic Regression

A classic pipeline using TF‑IDF features and a multinomial logistic regression classifier.

import pandas as pd
import os

test_text = pd.read_csv('dataset/test_text.csv')
train_all = pd.read_csv('dataset/train_all.csv')
example = pd.read_csv('example.csv')

# Convert categorical labels to numbers
category_mapping = {}
for k, v in enumerate(set(train_all['类别'].tolist())):
    category_mapping[v] = k
    print(k, v)

y_train = train_all['类别'].map(category_mapping)

import jieba

data = pd.concat([train_all[['id', '文本']], test_text[['id', '文本']]])

data['text'] = data['文本'].apply(lambda x: ' '.join(jieba.cut(x)))

test_text['text'] = test_text['文本'].apply(lambda x: ' '.join(jieba.cut(x)))
train_all['text'] = train_all['文本'].apply(lambda x: ' '.join(jieba.cut(x)))

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, auc, roc_auc_score
import joblib

vectorizer_word = TfidfVectorizer(
    max_features=800000,
    token_pattern=r"(?u)\b\w+\b",
    min_df=1,
    analyzer='word',
    ngram_range=(1, 3)
)
vectorizer_word = vectorizer_word.fit(data['text'])

tfidf_train = vectorizer_word.transform(train_all['text'])
tfidf_test = vectorizer_word.transform(test_text['text'])

clf = LogisticRegression(random_state=0, multi_class='multinomial')
clf.fit(tfidf_train, y_train)
joblib.dump(clf, 'clf.pkl')

model = joblib.load('clf.pkl')
y_pred_word = model.predict(tfidf_test)
reverse_mapping = {v: k for k, v in category_mapping.items()}
example['类别'] = y_pred_word
example['类别'] = example['类别'].map(reverse_mapping)
example.to_csv('res_0720_02.csv', header=True, index=False, encoding='UTF-8')

This approach achieved a macro F1 of about 0.48.

Baseline 2 – BERT

Fine‑tuning a Chinese BERT model for sequence classification.

import pandas as pd
import os
train_all = pd.read_csv('dataset/train_all.csv')
test_text = pd.read_csv('dataset/test_text.csv')
example = pd.read_csv('example.csv')

category_mapping = {}
for k, v in enumerate(set(train_all['类别'].tolist())):
    category_mapping[v] = k
train_all['label'] = train_all['类别'].map(category_mapping)

import torch
import random
from tqdm import tqdm
from transformers import BertTokenizer, BertForSequenceClassification, AdamW

num_classes = len(train_all['类别'].unique())

tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
model = BertForSequenceClassification.from_pretrained('bert-base-chinese', num_labels=num_classes)

class SentimentDataset(torch.utils.data.Dataset):
    def __init__(self, dataframe, tokenizer, max_length=128):
        self.dataframe = dataframe
        self.tokenizer = tokenizer
        self.max_length = max_length
    def __len__(self):
        return len(self.dataframe)
    def __getitem__(self, idx):
        text = self.dataframe.iloc[idx]['文本']
        label = self.dataframe.iloc[idx]['label']
        encoding = self.tokenizer(
            text,
            padding='max_length',
            truncation=True,
            max_length=self.max_length,
            return_tensors='pt'
        )
        return {
            'input_ids': encoding['input_ids'].flatten(),
            'attention_mask': encoding['attention_mask'].flatten(),
            'labels': torch.tensor(label, dtype=torch.long)
        }

dataset = SentimentDataset(train_all[['文本', 'label']], tokenizer)
train_size = int(0.95 * len(dataset))
val_size = len(dataset) - train_size
train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, val_size])

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=8, shuffle=True)
val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=8, shuffle=False)

optimizer = AdamW(model.parameters(), lr=5e-5)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

model.train()
for epoch in range(10):
    for batch in tqdm(train_loader, desc=f"Epoch {epoch+1}"):
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)
        optimizer.zero_grad()
        outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
        loss = outputs.loss
        loss.backward()
        optimizer.step()

model.eval()
total_eval_accuracy = 0
for batch in tqdm(val_loader, desc="Evaluating"):
    input_ids = batch['input_ids'].to(device)
    attention_mask = batch['attention_mask'].to(device)
    labels = batch['labels'].to(device)
    with torch.no_grad():
        outputs = model(input_ids, attention_mask=attention_mask)
    logits = outputs.logits
    preds = torch.argmax(logits, dim=1)
    accuracy = (preds == labels).float().mean()
    total_eval_accuracy += accuracy.item()
average_eval_accuracy = total_eval_accuracy / len(val_loader)
print("Validation Accuracy:", average_eval_accuracy)

def predict_sentiment(sentence):
    inputs = tokenizer(sentence, padding='max_length', truncation=True, max_length=128, return_tensors='pt').to(device)
    with torch.no_grad():
        outputs = model(**inputs)
    logits = outputs.logits
    preds = torch.argmax(logits, dim=1).item()
    return preds

preds_test = []
for text in test_text['文本'].tolist():
    preds_test.append(predict_sentiment(text))

reverse_mapping = {v: k for k, v in category_mapping.items()}
res = example.copy()
res['类别'] = preds_test
res['类别'] = res['类别'].map(reverse_mapping)
res.to_csv('res_0728_02.csv', header=True, index=False, encoding='UTF-8')

Baseline 3 – Large Model

Prompt‑based inference using the Qwen‑3‑0.6B model (or larger) via ModelScope.

import pandas as pd
import os

test_text = pd.read_csv('dataset/test_text.csv')
train_all = pd.read_csv('dataset/train_all.csv')
example = pd.read_csv('example.csv')

# Install and import ModelScope
!pip install modelscope
import modelscope
!modelscope download --model Qwen/Qwen3-0.6B

from modelscope import AutoModelForCausalLM, AutoTokenizer
import torch

path = "Qwen/Qwen3-0.6B"
tokenizer = AutoTokenizer.from_pretrained(path, torch_dtype=torch.float32)
model = AutoModelForCausalLM.from_pretrained(path, torch_dtype=torch.float32, device_map="auto")

prompt_template = '''
You are a professional text risk classifier. Classify the input text into one of the following 10 risk categories (output only the full category name, no explanations):
1.["政治敏感"]
2.["s情"]
3.["种族歧视"]
4.["地域歧视"]
5.["微侵犯(MA)"]
6.["犯罪"]
7.["基于文化背景的刻板印象(SCB)"]
8.["宗教迷信"]
9.["x侵犯(SO)"]
10.["基于外表的刻板印象(SA)"]

# Rules
- If multiple categories apply, choose the most prominent one.
- Always output a single category name.

{{TEXT}}
'''

labels = []
for id_, txt in zip(test_text['id'].tolist(), test_text['文本'].tolist()):
    prompt = prompt_template.replace('{{TEXT}}', txt)
    messages = [{"role": "user", "content": prompt}]
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=False)
    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
    generated_ids = model.generate(**model_inputs, max_new_tokens=32768)
    output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()
    # Simple parsing (ignore thinking tokens)
    content = tokenizer.decode(output_ids, skip_special_tokens=True).strip()
    labels.append(content)
    print("ID:", id_)
    print("Content:", txt)
    print("Result:", content)

res = example.copy()
res['类别'] = labels
res.to_csv('res_0720_01.csv', header=True, index=False, encoding='UTF-8')

The small Qwen‑3‑0.6B model performed poorly (macro F1 ≈ 0.0035). Fine‑tuning on the training set or using larger models such as DeepSeek can achieve scores around 90%.