crackrammer/ShieldBERT-Base-Chinese-Sensitive

1"""2推理脚本：使用训练好的模型预测文本是否包含敏感内容3"""4 5import os6import json7import argparse8 9import torch10from transformers import BertTokenizer, BertForSequenceClassification11 12 13class SensitiveWordPredictor:14    """敏感词预测器"""15 16    def __init__(self, model_path: str, device: str = None):17        if device:18            self.device = torch.device(device)19        else:20            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")21 22        # 加载配置23        config_path = os.path.join(model_path, "filter_config.json")24        with open(config_path, "r", encoding="utf-8") as f:25            self.config = json.load(f)26 27        self.label_map = self.config["label_map"]28        self.max_length = self.config.get("max_length", 128)29 30        # 加载模型和 tokenizer31        self.tokenizer = BertTokenizer.from_pretrained(model_path)32        self.model = BertForSequenceClassification.from_pretrained(model_path)33        self.model.to(self.device)34        self.model.eval()35 36    def predict(self, text: str) -> dict:37        """预测单条文本"""38        encoding = self.tokenizer(39            text,40            padding="max_length",41            truncation=True,42            max_length=self.max_length,43            return_tensors="pt",44        )45 46        input_ids = encoding["input_ids"].to(self.device)47        attention_mask = encoding["attention_mask"].to(self.device)48 49        with torch.no_grad():50            outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)51            probs = torch.softmax(outputs.logits, dim=1)52            pred_label = torch.argmax(probs, dim=1).item()53            confidence = probs[0][pred_label].item()54 55        return {56            "text": text,57            "label": pred_label,58            "label_name": self.label_map[str(pred_label)],59            "confidence": round(confidence, 4),60            "is_sensitive": pred_label == 1,61        }62 63    def predict_batch(self, texts: list[str], batch_size: int = 32) -> list[dict]:64        """批量预测"""65        results = []66        for i in range(0, len(texts), batch_size):67            batch_texts = texts[i : i + batch_size]68            encoding = self.tokenizer(69                batch_texts,70                padding="max_length",71                truncation=True,72                max_length=self.max_length,73                return_tensors="pt",74            )75 76            input_ids = encoding["input_ids"].to(self.device)77            attention_mask = encoding["attention_mask"].to(self.device)78 79            with torch.no_grad():80                outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)81                probs = torch.softmax(outputs.logits, dim=1)82                pred_labels = torch.argmax(probs, dim=1)83 84            for j, text in enumerate(batch_texts):85                label = pred_labels[j].item()86                conf = probs[j][label].item()87                results.append({88                    "text": text,89                    "label": label,90                    "label_name": self.label_map[str(label)],91                    "confidence": round(conf, 4),92                    "is_sensitive": label == 1,93                })94 95        return results96 97 98def main():99    parser = argparse.ArgumentParser(description="敏感词预测")100    parser.add_argument("--model_path", type=str, default="output/best_model")101    parser.add_argument("--text", type=str, help="要检测的文本（单条）")102    parser.add_argument("--file", type=str, help="要检测的文本文件（每行一条）")103    parser.add_argument("--device", type=str, default=None)104    args = parser.parse_args()105 106    predictor = SensitiveWordPredictor(args.model_path, args.device)107 108    if args.text:109        result = predictor.predict(args.text)110        status = "🔴 敏感" if result["is_sensitive"] else "🟢 正常"111        print(f"\n输入: {result['text']}")112        print(f"结果: {status}")113        print(f"标签: {result['label_name']} (label={result['label']})")114        print(f"置信度: {result['confidence']:.4f}")115 116    elif args.file:117        with open(args.file, "r", encoding="utf-8") as f:118            texts = [line.strip() for line in f if line.strip()]119 120        results = predictor.predict_batch(texts)121        print(f"\n{'='*70}")122        print(f"批量检测结果 (共 {len(results)} 条)")123        print(f"{'='*70}")124 125        sensitive_count = 0126        for r in results:127            status = "🔴 敏感" if r["is_sensitive"] else "🟢 正常"128            print(f"  {status} [{r['confidence']:.4f}] {r['text'][:50]}...")129            if r["is_sensitive"]:130                sensitive_count += 1131 132        print(f"\n统计: 正常 {len(results) - sensitive_count} 条, 敏感 {sensitive_count} 条")133 134    else:135        print("敏感词检测系统 - 交互模式")136        print("输入文本进行检测，输入 'quit' 退出\n")137        while True:138            text = input("请输入文本> ").strip()139            if text.lower() in ("quit", "exit", "q"):140                print("再见！")141                break142            if not text:143                continue144            result = predictor.predict(text)145            status = "🔴 敏感" if result["is_sensitive"] else "🟢 正常"146            print(f"  结果: {status} | 置信度: {result['confidence']:.4f}\n")147 148 149if __name__ == "__main__":150    main()151