Create model_a.py (#9)

- Create model_a.py (6d7ca919d7cf5f6f29304977e82a8aa97a15301e)


Co-authored-by: liujiaxiang <liujiaxiang@users.noreply.huggingface.co>

model_a.py

@@ -0,0 +1,14 @@
+# model_wrappers/model_a.py
+from transformers import pipeline
+
+# 中文 GPT2：Wenzhong
+model_a = pipeline("text-generation", model="IDEA-CCNL/Wenzhong-GPT2-110M", tokenizer="IDEA-CCNL/Wenzhong-GPT2-110M")
+
+def run_model_a(prompt: str) -> str:
+    output = model_a(prompt,
+                     max_length=100,
+                     do_sample=True,
+                     temperature=0.8,
+                     top_k=50,
+                     top_p=0.95)
+    return output[0]["generated_text"]