fix

utils/summary_utils.py

@@ -3,12 +3,11 @@ from transformers import pipeline
 
 from utils.tag_utils import filter_tags
 
-AiSummaryVersion = 3
-MinTagScore = 0.7
-summarization_pipeline = pipeline("summarization", model="Falconsai/text_summarization")
+AiSummaryVersion = 4
+summarization_pipeline = pipeline("summarization", model="csebuetnlp/mT5_multilingual_XLSum")
 en_translation_pipe = pipeline("translation", model="Helsinki-NLP/opus-mt-mul-en")
-tag_gen_pipe_1 = pipeline("text-classification", model="dima806/news-category-classifier-distilbert")
-tag_gen_pipe_2 = pipeline("text-classification", model="elozano/bert-base-cased-news-category")
+classification_pipe = pipeline("text-classification", model="Yueh-Huan/news-category-classification-distilbert")
+text_to_tags_pipe = pipeline('text2text-generation', model='models/text2tags')
 
 
 def summarize(id: str, text: str):
@@ -16,10 +15,11 @@ def summarize(id: str, text: str):
         return {
             "ver": AiSummaryVersion
         }
-    summary = get_summarization(text) if len(text) > 3000 else text
+    summary = get_summarization(text) if len(text) > 100 else text
     translated = get_en_translation(summary)
-    tags = get_tags(translated, id)
-    tags = filter_tags(tags)
+    tags1 = get_classification(translated)
+    tags2 = get_tags(translated)
+    tags = filter_tags(tags1 + tags2)
     tags = sorted(list(set(tags)))
 
     value = {
@@ -33,8 +33,7 @@ def summarize(id: str, text: str):
 
 def get_summarization(text: str):
     try:
-        # Max / Min number of words
-        result = summarization_pipeline(text, max_length=500, min_length=100, do_sample=False)
+        result = summarization_pipeline(text)
         return result[0]['summary_text'] if isinstance(result, list) else result['summary_text']
     except:
         return None
@@ -60,12 +59,25 @@ def is_english(text):
         return False
 
 
-def get_tags(text: str, id: str):
+def get_tags(text: str):
     if text is None:
         return []
     try:
-        tags1 = [tag['label'] for tag in tag_gen_pipe_1(text) if tag['score'] >= MinTagScore]
-        tags2 = [tag['label'] for tag in tag_gen_pipe_2(text) if tag['score'] >= MinTagScore]
-        return tags1 + tags2
+        result = text_to_tags_pipe(text)
+        tag_str = result[0]['generated_text'] if isinstance(result, list) else result['generated_text']
+        return [tag.strip() for tag in tag_str.split(',')]
+    except:
+        return []
+
+
+def get_classification(text: str):
+    if text is None:
+        return []
+    try:
+        result = classification_pipe(text)
+        if isinstance(result, list):
+            return [tag['label'].strip() for tag in result if tag['score'] > 0.75]
+        else:
+            return [result['label'].strip()] if result['score'] > 0.75 else []
     except:
         return []