CoolFace
Modelpublic

huiqian/tiny-sentiment-classifier

sourceHugging Facemitupdated 8mo agoView on Hugging Face
0likes
inference.py60 linesDownload Raw Back to root
1# inference.py (完整手动加载版本,绕过所有 Auto 机制)2 3import torch4from modeling_tinytransformer import TinyTransformerModel   # 直接导入你的模型类5from tokenization_tinytransformer import TinyTokenizer      # 你的 tokenizer6from configuration_tinytransformer import TinyTransformerConfig  # 你的 config 类7 8model_path = "./tiny-sentiment-model"9device = "cuda" if torch.cuda.is_available() else "cpu"10 11print("加载 tokenizer(手动)...")12tokenizer = TinyTokenizer()13 14print("加载 config(手动)...")15config = TinyTransformerConfig.from_pretrained(model_path)  # 用你的 config 类加载16 17print("加载模型(手动)...")18model = TinyTransformerModel.from_pretrained(model_path, config=config)19model.to(device)20model.eval()21 22print("模型加载完成!设备:", device)23 24# ----------------------- 预测函数(保持不变) -----------------------25def predict(text: str):26    inputs = tokenizer(27        text,28        padding=True,29        truncation=True,30        max_length=64,31        return_tensors="pt"32    )33    inputs = {k: v.to(device) for k, v in inputs.items()}34    35    with torch.no_grad():36        outputs = model(**inputs)37    38    logits = outputs["logits"] if isinstance(outputs, dict) else outputs.logits39    probs = torch.softmax(logits, dim=-1)[0]40    41    pred_label = torch.argmax(probs).item()42    confidence = probs[pred_label].item()43    44    label_map = {0: "负面", 1: "正面"}45    46    print(f"文本: {text}")47    print(f"预测: {label_map[pred_label]} (置信度: {confidence:.4f})")48    print(f"负面概率: {probs[0]:.4f} | 正面概率: {probs[1]:.4f}")49    print("-"*50)50 51# 测试52if __name__ == "__main__":53    test_texts = [54        "质量超级好,强烈推荐",55        "服务差劲,东西还贵",56        "一般般,没什么惊喜",57        "老板人好,下次还来"58    ]59    for text in test_texts:60        predict(text)