Files changed (1) hide show
  1. model_a.py +14 -0
model_a.py ADDED
@@ -0,0 +1,14 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # model_wrappers/model_a.py
2
+ from transformers import pipeline
3
+
4
+ # 中文 GPT2:Wenzhong
5
+ model_a = pipeline("text-generation", model="IDEA-CCNL/Wenzhong-GPT2-110M", tokenizer="IDEA-CCNL/Wenzhong-GPT2-110M")
6
+
7
+ def run_model_a(prompt: str) -> str:
8
+ output = model_a(prompt,
9
+ max_length=100,
10
+ do_sample=True,
11
+ temperature=0.8,
12
+ top_k=50,
13
+ top_p=0.95)
14
+ return output[0]["generated_text"]