Update predictor.py

predictor.py

@@ -1,155 +1,158 @@
-# predictor.py
-
-import torch
-import re
-import os
-import json
-import onnxruntime as ort
-from collections import defaultdict, Counter
-from difflib import SequenceMatcher
-from transformers import BertTokenizerFast
-
-
-class DialogueEvaluator:
-    def __init__(self, model_dir, keywords_path):
-        print("Initializing DialogueEvaluator...")
-        # 加载模型和tokenizer
-        self.tokenizer, self.model, self.id2label = self._load_model(model_dir)
-        print("✅ Model and Tokenizer loaded.")
-
-        # 从JSON文件加载关键词体系
-        with open(keywords_path, 'r', encoding='utf-8') as f:
-            self.academic_keywords = json.load(f)
-        print("✅ Keywords loaded.")
-
-        # 构建关键词正则表达式模式
-        self.keyword_patterns = self._build_keyword_patterns()
-        print("✅ Keyword patterns built.")
-
-        # 场景化停用词
-        self.scene_stopwords = r'^(嗯|啊|哦|呃|呐|哟)'
-        print("DialogueEvaluator initialized successfully.")
-
-    def _load_model(self, model_dir):
-        tokenizer = BertTokenizerFast.from_pretrained(model_dir)
-        sess_options = ort.SessionOptions()
-        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
-        model_path = os.path.join(model_dir, "model_quantized.onnx")
-
-        try:
-            model = ort.InferenceSession(model_path, sess_options, providers=providers)
-            print(f"ℹ️ ONNX Runtime using: {model.get_providers()[0]}")
-        except Exception as e:
-            print(f"❌ ONNX Runtime initialization failed: {str(e)}")
-            raise
-
-        with open(os.path.join(model_dir, "label_map.json"), 'r', encoding='utf-8') as f:
-            label_map = json.load(f)
-            id2label = {int(k): v for k, v in label_map['id2label'].items()}
-        return tokenizer, model, id2label
-
-    def _build_keyword_patterns(self):
-        patterns = {}
-        for scene, sentiment_dict in self.academic_keywords.items():
-            for sentiment, keywords in sentiment_dict.items():
-                for keyword in keywords:
-                    pattern = self._create_fuzzy_pattern(keyword)
-                    patterns[keyword] = {
-                        'pattern': pattern,
-                        'scene': scene,
-                        'sentiment': sentiment
-                    }
-        return patterns
-
-    def _create_fuzzy_pattern(self, keyword):
-        # 简单的模糊匹配，允许中间有一个任意字符
-        if len(keyword) <= 2:
-            return re.compile(re.escape(keyword))
-        pattern_str = re.escape(keyword[0]) + ''.join([f"{re.escape(c)}.?" for c in keyword[1:]])
-        return re.compile(pattern_str)
-
-    def _fuzzy_match_keywords(self, sentence):
-        matched_info = []
-        for keyword, info in self.keyword_patterns.items():
-            if info['pattern'].search(sentence):
-                # 简单处理否定情况
-                sentiment = info['sentiment']
-                if re.search(fr'(不|没有|无|否|缺乏|不足|不够){keyword}', sentence):
-                    sentiment = 'negative'
-
-                matched_info.append({
-                    'keyword': keyword,
-                    'scene': info['scene'],
-                    'sentiment': sentiment
-                })
-        return matched_info
-
-    def _clean_sentence(self, sentence):
-        sentence = re.sub(r'[^\w\s\u4e00-\u9fff，。；：、]', '', sentence)
-        sentence = re.sub(r'\s+', ' ', sentence).strip()
-        sentence = re.sub(self.scene_stopwords, '', sentence)
-        return sentence
-
-    def _extract_key_sentences(self, text):
-        sentences = re.split(r'[。！？；\n]', text)
-        key_sentences = []
-        for sent in sentences:
-            if len(sent) < 5: continue  # 忽略太短的句子
-            clean_sent = self._clean_sentence(sent)
-            if not clean_sent: continue
-            matched_info = self._fuzzy_match_keywords(clean_sent)
-            if matched_info:
-                key_sentences.append({
-                    'sentence': clean_sent,
-                    'matched_info': matched_info,
-                })
-        return key_sentences
-
-    def _predict_sentence(self, sentence):
-        inputs = self.tokenizer(
-            sentence, truncation=True, padding='max_length', max_length=128, return_tensors="np"
-        )
-        ort_inputs = {'input_ids': inputs['input_ids'], 'attention_mask': inputs['attention_mask']}
-        try:
-            outputs = self.model.run(None, ort_inputs)
-            logits = outputs[0]
-            probs = torch.softmax(torch.tensor(logits), dim=1)
-            pred_id = torch.argmax(probs).item()
-            return {
-                'label': self.id2label[pred_id],
-                'confidence': round(torch.max(probs).item(), 4)
-            }
-        except Exception as e:
-            print(f"❌ Inference failed for sentence: '{sentence}'. Error: {str(e)}")
-            return {'label': 'ERROR', 'confidence': 0.0}
-
-    def evaluate_full_text(self, text):
-        key_sentences_info = self._extract_key_sentences(text)
-        if not key_sentences_info:
-            return {'status': 'no_key_sentences', 'message': '未检测到包含评价关键词的有效句子。'}
-
-        processed_sentences = []
-        for sent_info in key_sentences_info:
-            prediction = self._predict_sentence(sent_info['sentence'])
-            sent_info.update(prediction)
-            processed_sentences.append(sent_info)
-
-        # --- 生成统计数据 ---
-        overall_stats = defaultdict(lambda: defaultdict(int))
-        all_labels = [sent['label'] for sent in processed_sentences]
-        overall_stats['total_sentences'] = len(processed_sentences)
-        overall_stats['label_distribution'] = dict(Counter(all_labels))
-        overall_stats['avg_confidence'] = round(
-            sum(s['confidence'] for s in processed_sentences) / len(processed_sentences),
-            4) if processed_sentences else 0
-
-        for sent in processed_sentences:
-            for info in sent['matched_info']:
-                overall_stats['scene_distribution'][info['scene']] += 1
-                overall_stats['sentiment_distribution'][info['sentiment']] += 1
-
-        return {
-            'status': 'success',
-            'overall_stats': dict(overall_stats),
-            'key_sentences': processed_sentences
-        }
+# predictor.py
+
+import torch
+import re
+import os
+import json
+from pathlib import Path
+import onnxruntime as ort
+from collections import defaultdict, Counter
+from difflib import SequenceMatcher
+from transformers import BertTokenizerFast
+
+
+class DialogueEvaluator:
+    def __init__(self, model_path, keywords_path):
+        print("Initializing DialogueEvaluator...")
+
+        # 加载模型和tokenizer
+        self.tokenizer, self.model, self.id2label = self._load_model(model_path)
+        print("✅ Model and Tokenizer loaded.")
+
+        # 加载关键词体系
+        with open(keywords_path, 'r', encoding='utf-8') as f:
+            self.academic_keywords = json.load(f)
+        print("✅ Keywords loaded.")
+
+        # 构建关键词正则
+        self.keyword_patterns = self._build_keyword_patterns()
+        print("✅ Keyword patterns built.")
+
+        # 场景化停用词
+        self.scene_stopwords = r'^(嗯|啊|哦|呃|呐|哟)'
+        print("DialogueEvaluator initialized successfully.")
+
+    def _load_model(self, model_path):
+        model_path = Path(model_path)
+        model_dir = model_path.parent
+
+        tokenizer = BertTokenizerFast.from_pretrained(model_dir)
+
+        sess_options = ort.SessionOptions()
+        providers = ['CPUExecutionProvider']  # Hugging Face Spaces 通常无 GPU
+        model = ort.InferenceSession(str(model_path), sess_options, providers=providers)
+        print(f"ℹ️ ONNX Runtime using: {model.get_providers()[0]}")
+
+        label_map_path = model_dir / "label_map.json"
+        if not label_map_path.exists():
+            raise FileNotFoundError(f"Missing label_map.json at: {label_map_path}")
+        
+        with open(label_map_path, 'r', encoding='utf-8') as f:
+            label_map = json.load(f)
+            id2label = {int(k): v for k, v in label_map['id2label'].items()}
+        return tokenizer, model, id2label
+
+    def _build_keyword_patterns(self):
+        patterns = {}
+        for scene, sentiment_dict in self.academic_keywords.items():
+            for sentiment, keywords in sentiment_dict.items():
+                for keyword in keywords:
+                    pattern = self._create_fuzzy_pattern(keyword)
+                    patterns[keyword] = {
+                        'pattern': pattern,
+                        'scene': scene,
+                        'sentiment': sentiment
+                    }
+        return patterns
+
+    def _create_fuzzy_pattern(self, keyword):
+        if len(keyword) <= 2:
+            return re.compile(re.escape(keyword))
+        pattern_str = re.escape(keyword[0]) + ''.join([f"{re.escape(c)}.?" for c in keyword[1:]])
+        return re.compile(pattern_str)
+
+    def _fuzzy_match_keywords(self, sentence):
+        matched_info = []
+        for keyword, info in self.keyword_patterns.items():
+            if info['pattern'].search(sentence):
+                sentiment = info['sentiment']
+                if re.search(fr'(不|没有|无|否|缺乏|不足|不够){keyword}', sentence):
+                    sentiment = 'negative'
+
+                matched_info.append({
+                    'keyword': keyword,
+                    'scene': info['scene'],
+                    'sentiment': sentiment
+                })
+        return matched_info
+
+    def _clean_sentence(self, sentence):
+        sentence = re.sub(r'[^\w\s\u4e00-\u9fff，。；：、]', '', sentence)
+        sentence = re.sub(r'\s+', ' ', sentence).strip()
+        sentence = re.sub(self.scene_stopwords, '', sentence)
+        return sentence
+
+    def _extract_key_sentences(self, text):
+        sentences = re.split(r'[。！？；\n]', text)
+        key_sentences = []
+        for sent in sentences:
+            if len(sent) < 5:
+                continue
+            clean_sent = self._clean_sentence(sent)
+            if not clean_sent:
+                continue
+            matched_info = self._fuzzy_match_keywords(clean_sent)
+            if matched_info:
+                key_sentences.append({
+                    'sentence': clean_sent,
+                    'matched_info': matched_info,
+                })
+        return key_sentences
+
+    def _predict_sentence(self, sentence):
+        inputs = self.tokenizer(
+            sentence, truncation=True, padding='max_length', max_length=128, return_tensors="np"
+        )
+        ort_inputs = {'input_ids': inputs['input_ids'], 'attention_mask': inputs['attention_mask']}
+        try:
+            outputs = self.model.run(None, ort_inputs)
+            logits = outputs[0]
+            probs = torch.softmax(torch.tensor(logits), dim=1)
+            pred_id = torch.argmax(probs).item()
+            return {
+                'label': self.id2label[pred_id],
+                'confidence': round(torch.max(probs).item(), 4)
+            }
+        except Exception as e:
+            print(f"❌ Inference failed for sentence: '{sentence}'. Error: {str(e)}")
+            return {'label': 'ERROR', 'confidence': 0.0}
+
+    def evaluate_full_text(self, text):
+        key_sentences_info = self._extract_key_sentences(text)
+        if not key_sentences_info:
+            return {'status': 'no_key_sentences', 'message': '未检测到包含评价关键词的有效句子。'}
+
+        processed_sentences = []
+        for sent_info in key_sentences_info:
+            prediction = self._predict_sentence(sent_info['sentence'])
+            sent_info.update(prediction)
+            processed_sentences.append(sent_info)
+
+        overall_stats = defaultdict(lambda: defaultdict(int))
+        all_labels = [sent['label'] for sent in processed_sentences]
+        overall_stats['total_sentences'] = len(processed_sentences)
+        overall_stats['label_distribution'] = dict(Counter(all_labels))
+        overall_stats['avg_confidence'] = round(
+            sum(s['confidence'] for s in processed_sentences) / len(processed_sentences),
+            4) if processed_sentences else 0
+
+        for sent in processed_sentences:
+            for info in sent['matched_info']:
+                overall_stats['scene_distribution'][info['scene']] += 1
+                overall_stats['sentiment_distribution'][info['sentiment']] += 1
+
+        return {
+            'status': 'success',
+            'overall_stats': dict(overall_stats),
+            'key_sentences': processed_sentences
+        }