Update model_wrappers/model_b.py

#7
Files changed (1) hide show
  1. model_wrappers/model_b.py +3 -1
model_wrappers/model_b.py CHANGED
@@ -1,6 +1,8 @@
 
1
  from transformers import pipeline
2
 
3
- model_b = pipeline("text-generation", model="distilgpt2")
 
4
 
5
  def run_model_b(prompt: str) -> str:
6
  output = model_b(prompt,
 
1
+ # model_wrappers/model_b.py
2
  from transformers import pipeline
3
 
4
+ # 中文 GPT 模型:Mengzi GPT
5
+ model_b = pipeline("text-generation", model="Langboat/mengzi-gpt-base", tokenizer="Langboat/mengzi-gpt-base")
6
 
7
  def run_model_b(prompt: str) -> str:
8
  output = model_b(prompt,