class for model single

ebart.py

@@ -1,46 +1,66 @@
 import spaces
 import torch
-from transformers import PegasusForConditionalGeneration, PegasusTokenizer
-#from tokenizers_pegasus import PegasusTokenizer
+from transformers import PegasusForConditionalGeneration
+from tokenizers_pegasus import PegasusTokenizer
+
+class PegasusSummarizer:
+    _instance = None
+
+    def __new__(cls, *args, **kwargs):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+            cls._instance._init_model()
+        return cls._instance
+
+    def _init_model(self):
+        # 加载标记器和模型
+        model_name = "IDEA-CCNL/Randeng-Pegasus-523M-Summary-Chinese-V1"
+        self.model = PegasusForConditionalGeneration.from_pretrained(model_name)
+        self.tokenizer = PegasusTokenizer.from_pretrained(model_name)
+
+        # 将模型移动到GPU
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model.to(self.device)
+
+    def generate_summary(self, text, max_length=180, min_length=64):
+        # 进行标记化并将输入数据移动到GPU
+        inputs = self.tokenizer(text, max_length=1024, truncation=True, return_tensors="pt").to(self.device)
+
+        # 生成摘要
+        summary_ids = self.model.generate(
+            inputs["input_ids"],
+            max_length=max_length,
+            min_length=min_length,
+            num_beams=4,
+            early_stopping=True,
+            temperature=0.7,
+            top_k=50,
+            top_p=0.9,
+            repetition_penalty=2.0,
+            length_penalty=1.0,
+            no_repeat_ngram_size=3,
+            num_return_sequences=1,
+            do_sample=True
+        )
+
+        # 解码并返回摘要
+        clean_summary = self.tokenizer.batch_decode(summary_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)[0]
+
+        # 处理并过滤掉不需要的特殊标记
+        special_tokens = ['<pad>', '<unk>', '</s>']
+        for token in special_tokens:
+            clean_summary = clean_summary.replace(token, '')
+
+        return clean_summary
 
 @spaces.GPU
 def generate_summary(text, max_length=180, min_length=64):
-    # 加载标记器和模型
-    model = PegasusForConditionalGeneration.from_pretrained("IDEA-CCNL/Randeng-Pegasus-523M-Summary-Chinese-V1")
-    tokenizer = PegasusTokenizer.from_pretrained("IDEA-CCNL/Randeng-Pegasus-523M-Summary-Chinese-V1")
-
-    # 将模型移动到GPU
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model.to(device)
-
-    # 进行标记化并将输入数据移动到GPU
-    inputs = tokenizer(text, max_length=1024, truncation=True, return_tensors="pt").to(device)
-
-    # 生成摘要
-    summary_ids = model.generate(
-        inputs["input_ids"],
-        max_length=max_length,
-        min_length=min_length,
-        num_beams=4,
-        early_stopping=True
-    )
-
-    # 解码并返回摘要
-    clean_summary = tokenizer.batch_decode(summary_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)[0]
-    return clean_summary
+    summarizer = PegasusSummarizer()
+    return summarizer.generate_summary(text, max_length, min_length)
 
 if __name__ == "__main__":
     text = (
-        "2023年3月16日我们从黑龙江大庆来到湖南长沙市长沙县福中路77号湖南省富达日化有限公司，"
-        "其宣传的特殊配方洗衣液比立白和蓝月亮的去污效果要好很多，还有油污净自称中国去污第一名，"
-        "做了一些去除废机油的实验，油污净清洗废机油，洗洗液去除废机油，洗完后直接排入城市管网的下水池，"
-        "每天都在进行相关测试，废机油属于危险废物，严重危害公共环境，请湖南环保局对其污染环境进行查处。"
-        "其公司宣传材料存在大量虚假宣传，夸大其词，感觉就是个传销组织，说其公司有妆字号资质，药字号资质，"
-        "全国工业产品餐具用洗涤剂资质，声称其设备是纯净水设备，是否有涉水批件，是否有消字号证件，"
-        "其消字号所有产品是否都进行备案和匹配的检测报告，请湖南市场监督管理局对其进行查处，"
-        "1997年到现在坑害全国百姓加盟其公司，请湖南商务局查处其是否具有特许经营资质，自称每年营业额1亿元多元，"
-        "从97年坑害到23年大量的客户没有开局相应的发票，存在重大偷税漏税嫌疑，请湖南税务机关对其进行查处！"
-        "还有其出口的设备，渠道是否正规，是白关，灰关，还是黑关，请湖南海关相关部门对其进行查处。"
+        "东四路西侧之前有划分免费停车位，为什么后面被撤销，而道路东侧有划分免费停车车位，附近小区车位紧张导致很难找到停车位。"
     )
     summary = generate_summary(text, max_length=128, min_length=64)
     print(summary)