优化情感得分计算逻辑，增加长文本处理功能，改进模型得分组合方式，添加 nltk 依赖

blkeras.py

@@ -96,7 +96,7 @@ def ensure_fixed_shape(data, shape, variable_name=""):
 
 def predict(text: str, stock_codes: list):
     from tensorflow.keras.preprocessing.sequence import pad_sequences # type: ignore
-    from preprocess import get_document_vector, get_stock_info, preprocessing_entry, process_entities, process_pos_tags, processing_entry
+    from preprocess import get_document_vector, get_stock_info, process_entities, process_pos_tags, processing_entry
 
     try:
 
@@ -110,10 +110,8 @@ def predict(text: str, stock_codes: list):
         #print(f"predict() Input text: {input_text}")
 
         # 使用预处理函数处理文本
-        processed_entry = processing_entry(input_text)
-
         # 解包 processed_entry 中的各个值
-        lemmatized_entry, pos_tag, ner, _ , sentiment_score = processed_entry
+        lemmatized_entry, pos_tag, ner, _ , sentiment_score = processing_entry(input_text)
 
         # 分别打印每个变量，便于调试
         #print("Lemmatized Entry:", lemmatized_entry)

preprocess.py

@@ -223,72 +223,121 @@ def get_document_vector(words, model = word2vec_model):
 
 
 # 函数：获取情感得分
-def get_sentiment_score(text):
+def process_long_text(text, tokenizer, max_length=512):
+    """
+    将长文本分段并保持句子完整性
+    """
+    import nltk
     try:
-        import torch
-
-        # 获取第一个模型的结果 (ProsusAI/finbert)
-        # result_one = sentiment_analyzer_one(text, truncation=True, max_length=512)[0]
-        # 获取模型输出
-        with torch.no_grad():
-            outputs_one = sa_model_one(**tokenizer_one(text, return_tensors="pt", truncation=False))
-            predictions_one = torch.nn.functional.softmax(outputs_one.logits, dim=-1)
-
-
-            outputs_two = sa_model_two(**tokenizer_two(text, return_tensors="pt", truncation=False))
-            predictions_two = torch.nn.functional.softmax(outputs_two.logits, dim=-1)
-
-        # 获取所有标签的概率
-        scores_one = predictions_one[0].tolist()
-        scores_two = predictions_two[0].tolist()
-
-        # 获取标签映射
-        # labels_one = sa_model_one.config.id2label
-        # labels_two = sa_model_two.config.id2label
-
-        # 打印所有标签的概率
-        score_one_positive = scores_one[0]
-        score_one_negative = scores_one[1]
-        score_one_neutral = scores_one[2]
-
-
-        final_score_one = 0.0
-        final_score_one += score_one_positive
-        final_score_one -= score_one_negative
-        if score_one_positive > score_one_negative:
-            final_score_one += score_one_neutral
-        else:
-            final_score_one -= score_one_neutral
-
-        final_score_one = max(-1.0, min(1.0, final_score_one))
-
-        score_two_neutral = scores_two[0]
-        score_two_positive = scores_two[1]
-        score_two_negative = scores_two[2]
+        nltk.data.find('tokenizers/punkt')
+    except LookupError:
+        nltk.download('punkt')
 
-        final_score_two = 0.0
-        final_score_two += score_two_positive
-        final_score_two -= score_two_negative
-        if score_two_positive > score_two_negative:
-            final_score_two += score_two_neutral
+    try:
+        nltk.data.find('tokenizers/punkt_tab')
+    except LookupError:
+        nltk.download('punkt_tab')
+        
+        
+    sentences = nltk.sent_tokenize(text)
+    segments = []
+    current_segment = ""
+    
+    for sentence in sentences:
+        print(f"Processing sentence: {sentence}")
+        # 检查添加当前句子后是否会超过最大长度
+        test_segment = current_segment + " " + sentence if current_segment else sentence
+        if len(tokenizer.tokenize(test_segment)) > max_length:
+            if current_segment:
+                segments.append(current_segment.strip())
+            current_segment = sentence
         else:
-            final_score_two -= score_two_neutral
+            current_segment = test_segment
+            
+    # 添加最后一个段落
+    if current_segment:
+        segments.append(current_segment.strip())
         
+    return segments
 
-        # 将两个模型的得分组合（加权平均）
-        final_score = np.average([final_score_one, final_score_two], weights=[0.3, 0.7])
+def get_sentiment_score(text):
+    try:
+        import torch
         
-        # 确保最终得分在 [-1, 1] 范围内
-        final_score = np.clip(final_score, -1.0, 1.0)
+        # 将长文本分段
+        segments_one = process_long_text(text, tokenizer_one)
+        segments_two = process_long_text(text, tokenizer_two)
         
-        return final_score
+        final_scores_one = []
+        final_scores_two = []
+        weights_one = []
+        weights_two = []
+
+        # 处理每个段落 - 模型一
+        for segment in segments_one:
+            with torch.no_grad():
+                inputs = tokenizer_one(segment, return_tensors="pt", truncation=True, max_length=512)
+                outputs = sa_model_one(**inputs)
+                predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
+                
+                scores = predictions[0].tolist()
+                score_positive = scores[0]
+                score_negative = scores[1]
+                score_neutral = scores[2]
+                
+                segment_score = 0.0
+                segment_score += score_positive
+                segment_score -= score_negative
+                if score_positive > score_negative:
+                    segment_score += score_neutral
+                else:
+                    segment_score -= score_neutral
+                
+                final_scores_one.append(np.clip(segment_score, -1.0, 1.0))
+                weights_one.append(len(tokenizer_one.tokenize(segment)))
+
+        # 处理每个段落 - 模型二
+        for segment in segments_two:
+            with torch.no_grad():
+                inputs = tokenizer_two(segment, return_tensors="pt", truncation=True, max_length=512)
+                outputs = sa_model_two(**inputs)
+                predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
+                
+                scores = predictions[0].tolist()
+                score_neutral = scores[0]
+                score_positive = scores[1]
+                score_negative = scores[2]
+                
+                segment_score = 0.0
+                segment_score += score_positive
+                segment_score -= score_negative
+                if score_positive > score_negative:
+                    segment_score += score_neutral
+                else:
+                    segment_score -= score_neutral
+                
+                final_scores_two.append(np.clip(segment_score, -1.0, 1.0))
+                weights_two.append(len(tokenizer_two.tokenize(segment)))
+
+        # 加权平均
+        if final_scores_one:
+            final_score_one = np.average(final_scores_one, weights=weights_one)
+        else:
+            final_score_one = 0.0
+            
+        if final_scores_two:
+            final_score_two = np.average(final_scores_two, weights=weights_two)
+        else:
+            final_score_two = 0.0
+
+        # 组合两个模型的结果
+        final_score = np.average([final_score_one, final_score_two], weights=[0.3, 0.7])
+        return np.clip(final_score, -1.0, 1.0)
 
     except Exception as e:
         print(f"Error in get_sentiment_score for text: {text[:50]}... Error: {str(e)}")
         traceback.print_exc()
         return 0.0
-    
-
 
 
 def get_stock_info(stock_code: str, history_days=30):

requirements.txt

@@ -16,4 +16,5 @@ yfinance==0.2.47
 jsonpath==0.82.2
 tensorflow==2.16.2
 pydantic==2.9.2
-pydantic_core==2.23.4
+pydantic_core==2.23.4
+nltk