Upload inference.py with huggingface_hub

inference.py

@@ -0,0 +1,187 @@
+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)