inference.py

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)