User Intent Recognition in Zhihu Scenarios

from __future__ import annotations

import json
import random
import zipfile
from pathlib import Path
from typing import Any

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
from transformers import AutoModel, AutoTokenizer, get_linear_schedule_with_warmup
from tqdm import tqdm
import random

seed = 42

random.seed(seed)                  # Python built-in random
np.random.seed(seed)               # NumPy
torch.manual_seed(seed)            # PyTorch (CPU)
torch.cuda.manual_seed(seed)       # PyTorch (single GPU)
torch.cuda.manual_seed_all(seed)   # PyTorch (all GPUs)

# Ensures deterministic behavior
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False



DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"DEVICE = {DEVICE}")
PATH = '/bohr/train-a3ld/v1/'
BASE_MODEL_NAME = "bert-base-chinese"
# MODEL_DIR = PATH + "models--bert-base-chinese/8f23c25b06e129b6c986331a13d8d025a92cf0ea"
MODEL_DIR = PATH + "bert-base-chinese"
TRAIN_PATH = Path(PATH + "/train.jsonl")


MAX_LEN = 128
BATCH_SIZE = 16
EPOCHS = 3
LEARNING_RATE = 2e-5
SEED = 42


def set_seed(s):
    random.seed(s); np.random.seed(s); torch.manual_seed(s); torch.cuda.manual_seed_all(s)


set_seed(SEED)

INTENT_LABELS = [
    "游戏技巧", "游戏角色信息", "周边饭馆", "查找高端酒店", "查找地址",  # game techniques, game character information, nearby restaurants, finding high-end hotels, finding addresses
    "健康知识", "查找医疗信息", "美容化妆技巧", "美食烹饪技巧", "软件开发问题",  # health knowledge, finding medical information, beauty and make-up techniques, cooking techniques, software development questions
    "软件使用问题", "查找小说剧情", "查找小说角色信息", "查找营业时间",  # software usage questions, finding novel plots, finding novel character information, finding opening hours
    "法律法条解释", "法律问题咨询",  # interpretation of legal provisions, legal consultation
]
LABEL_TO_ID = {l: i for i, l in enumerate(INTENT_LABELS)}
ID_TO_LABEL = {i: l for i, l in enumerate(INTENT_LABELS)}

def extract_user_utterance(text: str) -> str:
    user_utts = []
    for line in text.strip().split("\n"):
        if line.strip().startswith("usr:"):
            user_utts.append(line.strip()[4:].strip())
    return user_utts[-1] if user_utts else ""


def preprocess_input_text(text: Any) -> str:
    if text is None:
        return ""
    s = str(text).strip()
    if not s:
        return ""
    if "usr:" in s:
        u = extract_user_utterance(s)
        if u:
            return u
    return s


def load_jsonl(path):
    return [json.loads(l) for l in path.read_text(encoding="utf-8").splitlines() if l.strip()]


train_records = load_jsonl(TRAIN_PATH)
train_rows = []
for r in train_records:
    text = preprocess_input_text(r.get("input"))
    lab = str(r.get("label", "")).strip()
    if text and lab in LABEL_TO_ID:
        train_rows.append({"text": text, "label_id": LABEL_TO_ID[lab]})
print(f"Training samples: {len(train_rows)}")

from transformers import BertTokenizer, BertModel
class IntentDataset(Dataset):
    def __init__(self, rows, tokenizer, max_len, with_label=True):
        self.rows = rows; self.tokenizer = tokenizer; self.max_len = max_len; self.with_label = with_label

    def __len__(self):
        return len(self.rows)

    def __getitem__(self, idx):
        row = self.rows[idx]
        enc = self.tokenizer(row["text"], add_special_tokens=True, truncation=True,
                             max_length=self.max_len, padding="max_length", return_tensors="pt")
        item = {"input_ids": enc["input_ids"].flatten(), "attention_mask": enc["attention_mask"].flatten()}
        if self.with_label:
            item["labels"] = torch.tensor(row["label_id"], dtype=torch.long)
        return item


class BertIntentClassifier(nn.Module):
    def __init__(self, model_name, num_classes):
        super().__init__()
        self.bert = BertModel.from_pretrained(model_name)
        self.dropout = nn.Dropout(0.3)
        self.classifier = nn.Linear(self.bert.config.hidden_size, num_classes)

    def forward(self, input_ids, attention_mask):
        out = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        if isinstance(out, tuple):
            pooled = out[0][:, 0]
        else:
            pooled = getattr(out, "pooler_output", None)
            if pooled is None:
                pooled = out.last_hidden_state[:, 0]
        return self.classifier(self.dropout(pooled))


tokenizer = BertTokenizer.from_pretrained(MODEL_DIR)
model = BertIntentClassifier(MODEL_DIR, len(INTENT_LABELS)).to(DEVICE)

train_loader = DataLoader(IntentDataset(train_rows, tokenizer, MAX_LEN), batch_size=BATCH_SIZE, shuffle=True)
optim = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)
total_steps = max(1, len(train_loader) * EPOCHS)
warmup_steps = max(1, int(total_steps * 0.1))
sched = get_linear_schedule_with_warmup(optim, num_warmup_steps=warmup_steps, num_training_steps=total_steps)

for ep in range(EPOCHS):
    print(f"epoch {ep+1}/{EPOCHS}")
    model.train()
    for batch in tqdm(train_loader, desc="train", leave=False):
        optim.zero_grad()
        x = batch["input_ids"].to(DEVICE); m = batch["attention_mask"].to(DEVICE); y = batch["labels"].to(DEVICE)
        loss = nn.CrossEntropyLoss()(model(x, m), y)
        loss.backward(); optim.step(); sched.step()
print("training done.")

import os

if os.environ.get('DATA_PATH'):
    DATA_PATH = os.environ.get("DATA_PATH") + "/"
else:
    print("When the baseline runs, the test set cannot be read, so this error message appears; this is normal.")
    print("When baseline is running, this error message will appear because the test set cannot be read, which is a normal phenomenon.")

TEST_A_PATH = Path(DATA_PATH + "/val.jsonl")
TEST_B_PATH = Path(DATA_PATH + "/test.jsonl")


@torch.no_grad()
def predict_to_file(test_path, out_path):
    model.eval()
    records = load_jsonl(test_path)
    rows = [{"text": preprocess_input_text(r.get("input")) or ""} for r in records]
    loader = DataLoader(IntentDataset(rows, tokenizer, MAX_LEN, with_label=False), batch_size=BATCH_SIZE*4, shuffle=False)
    preds = []
    for batch in tqdm(loader, desc="predict", leave=False):
        x = batch["input_ids"].to(DEVICE); m = batch["attention_mask"].to(DEVICE)
        preds.extend(torch.argmax(model(x, m), dim=1).cpu().tolist())
    assert len(preds) == len(records)
    with Path(out_path).open("w", encoding="utf-8") as fh:
        for p in preds:
            fh.write(json.dumps({"label": ID_TO_LABEL[p]}, ensure_ascii=False) + "\n")
    print(f"wrote {len(preds)} -> {out_path}")


predict_to_file(TEST_A_PATH, Path("submission_val.jsonl"))
predict_to_file(TEST_B_PATH, Path("submission_test.jsonl"))

SUBMISSION_ZIP = Path("submission.zip")
with zipfile.ZipFile(SUBMISSION_ZIP, "w", zipfile.ZIP_DEFLATED) as zf:
    zf.write("submission_val.jsonl")
    zf.write("submission_test.jsonl")
print(f"wrote {SUBMISSION_ZIP} ({SUBMISSION_ZIP.stat().st_size} bytes)")

'''

import os
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# Expand the path of the user's home directory
model_path = os.path.expanduser("~/Qwen3-4B-Instruct")

# Detect the available device
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

# Load the tokeniser and the model (with automatically inferred precision; bfloat16 is used automatically if a GPU is available)
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,
    device_map="auto",
    trust_remote_code=True
)

# Define the dialogue messages (following Qwen's chat template)
messages = [
    {"role": "system", "content": "你是一个有用的助手。"},  # "You are a helpful assistant."
    {"role": "user", "content": "你好，请介绍一下你自己。"}  # "Hello, please introduce yourself."
]

# Apply the chat template to build the input
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True
)

# Encode the input and move it to the model's device
model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

# Generation parameters for the reply (adjust as needed)
generated_ids = model.generate(
    **model_inputs,
    max_new_tokens=512,
    do_sample=True,
    temperature=0.7,
    top_p=0.8,
    repetition_penalty=1.05
)

# Extract only the generated part (remove the original input)
generated_ids = [
    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
]

# Decode the reply
response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
print("\nModel reply:")
print(response)
'''