import os import json import torch import torch.nn as nn from transformers import AutoTokenizer, AutoModel import argparse from loguru import logger from tqdm import tqdm from torch.utils.data import DataLoader, Dataset class GlobalMeanPoolClassifier(nn.Module): def __init__(self, base_model, num_labels=2, dropout=0.0): super().__init__() self.base_model = base_model if hasattr(base_model.config, "text_config"): hidden_size = base_model.config.text_config.hidden_size else: hidden_size = base_model.config.hidden_size self.classifier = nn.Sequential( nn.LayerNorm(hidden_size), nn.Linear(hidden_size, 1024), nn.SiLU(), nn.Dropout(dropout), nn.Linear(1024, num_labels), ) def forward(self, input_ids, attention_mask, **kwargs): outputs = self.base_model( input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True, ) # Layer 19 for pure local features features = outputs.hidden_states[19] batch_size = input_ids.size(0) pooled_features = [] for i in range(batch_size): mask = attention_mask[i].bool() # Global Mean Pooling on non-pad tokens global_feat = features[i, mask, :].mean(dim=0) pooled_features.append(global_feat) pooled_features = torch.stack(pooled_features) logits = self.classifier(pooled_features) return logits PROMPT_TEMPLATE = ( "任务：论坛新词发现。\n" "根据候选词及其上下文，判断该词是否具有稳定独立的语义。\n" "候选词：{word}\n" "上下文：{context}" ) class InferenceDataset(Dataset): def __init__(self, items, tokenizer): self.items = items self.tokenizer = tokenizer def __len__(self): return len(self.items) def __getitem__(self, idx): item = self.items[idx] text = PROMPT_TEMPLATE.format(word=item["word"], context=item["context"]) return { "word": item["word"], "text": text } def collate_fn(batch, tokenizer): words = [b["word"] for b in batch] texts = [b["text"] for b in batch] encodings = tokenizer(texts, padding=True, truncation=True, max_length=256, return_tensors="pt") return words, encodings def run_inference(input_path, model_dir, output_dir, batch_size=32): device = "cuda" if torch.cuda.is_available() else "cpu" logger.info(f"Using device: {device}") logger.info(f"Loading Base Model and Tokenizer from {model_dir}...") tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token base_model = AutoModel.from_pretrained( model_dir, trust_remote_code=True, torch_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16, device_map="auto" ) logger.info("Initializing GlobalMeanPoolClassifier...") model = GlobalMeanPoolClassifier(base_model, num_labels=2) head_path = os.path.join(model_dir, "classifier_head.pt") if os.path.exists(head_path): logger.info("Loading trained classifier head weights...") model.classifier.load_state_dict(torch.load(head_path, map_location="cpu")) else: logger.warning(f"No classifier_head.pt found at {head_path}! Model will use random head weights.") model.to(device) model.eval() logger.info(f"Reading input data from {input_path}...") with open(input_path, "r", encoding="utf-8") as f: data = json.load(f) # Flatten data for batching items = [] for word, info in data.items(): contexts = info.get("contexts", []) if not contexts: continue # Take up to 5 contexts per word to speed up and avoid bias for ctx in contexts[:5]: items.append({"word": word, "context": ctx}) dataset = InferenceDataset(items, tokenizer) dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, collate_fn=lambda b: collate_fn(b, tokenizer)) logger.info("Running inference...") word_scores = {} with torch.no_grad(): for words, encodings in tqdm(dataloader): input_ids = encodings["input_ids"].to(device) attention_mask = encodings["attention_mask"].to(device) logits = model(input_ids=input_ids, attention_mask=attention_mask) probs = torch.softmax(logits, dim=-1) # Prob for class 0 (Accept) accept_probs = probs[:, 0].cpu().numpy() for word, prob in zip(words, accept_probs): if word not in word_scores: word_scores[word] = [] word_scores[word].append(float(prob)) # Aggregate scores (Mean probability across all contexts) accepted = {} rejected = {} logger.info("Aggregating results...") for word, probs in word_scores.items(): mean_prob = sum(probs) / len(probs) result_entry = { "score": round(mean_prob, 4), "contexts_analyzed": len(probs) } # Threshold at 0.5 if mean_prob >= 0.5: accepted[word] = result_entry else: rejected[word] = result_entry os.makedirs(output_dir, exist_ok=True) accept_path = os.path.join(output_dir, "accepted_words.json") reject_path = os.path.join(output_dir, "rejected_words.json") with open(accept_path, "w", encoding="utf-8") as f: json.dump(accepted, f, ensure_ascii=False, indent=2) with open(reject_path, "w", encoding="utf-8") as f: json.dump(rejected, f, ensure_ascii=False, indent=2) logger.info(f"Done! Evaluated {len(word_scores)} unique words.") logger.info(f"Accepted: {len(accepted)} words -> {accept_path}") logger.info(f"Rejected: {len(rejected)} words -> {reject_path}") if __name__ == "__main__": parser = argparse.ArgumentParser(description="Batch inference for BBS New Words Discovery") parser.add_argument("--input", type=str, required=True, help="Path to input JSON file (e.g., data.json)") parser.add_argument("--model_dir", type=str, default="./final_production_model", help="Path to merged production model") parser.add_argument("--output_dir", type=str, default="./inference_results", help="Directory to save accepted/rejected files") parser.add_argument("--batch_size", type=int, default=32, help="Batch size for inference") args = parser.parse_args() run_inference(args.input, args.model_dir, args.output_dir, args.batch_size)