change model to randeng

ebart.py

@@ -1,30 +1,33 @@
 import spaces
 import torch
-from transformers import PegasusForConditionalGeneration
-# 从 Fengshenbang-LM 下载 tokenizers_pegasus.py 和其他 Python 脚本
-from tokenizers_pegasus import PegasusTokenizer
+from transformers import PegasusForConditionalGeneration, PegasusTokenizer
+#from tokenizers_pegasus import PegasusTokenizer
 
 @spaces.GPU
 def generate_summary(text, max_length=180, min_length=64):
     # 加载标记器和模型
-    model = PegasusForConditionalGeneration.from_pretrained("IDEA-CCNL/Randeng-Pegasus-238M-Summary-Chinese")
-    tokenizer = PegasusTokenizer.from_pretrained("IDEA-CCNL/Randeng-Pegasus-238M-Summary-Chinese")
-            
+    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)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
 
-    # 进行标记化
-    inputs = tokenizer(text, max_length=1024, return_tensors="pt")#.to(device)
+    # 进行标记化并将输入数据移动到GPU
+    inputs = tokenizer(text, max_length=1024, truncation=True, return_tensors="pt").to(device)
 
     # 生成摘要
     summary_ids = model.generate(
-            inputs["input_ids"]
-        )
-    clean_summary = tokenizer.batch_decode(summary_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
-    print(clean_summary)
-    
+        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
 
 if __name__ == "__main__":
     text = (
@@ -39,4 +42,5 @@ if __name__ == "__main__":
         "从97年坑害到23年大量的客户没有开局相应的发票，存在重大偷税漏税嫌疑，请湖南税务机关对其进行查处！"
         "还有其出口的设备，渠道是否正规，是白关，灰关，还是黑关，请湖南海关相关部门对其进行查处。"
     )
-    generate_summary(text, max_length=128, min_length=64)
+    summary = generate_summary(text, max_length=128, min_length=64)
+    print(summary)