Tong commited on
Commit
f07a852
1 Parent(s): 36ce433

update code

Browse files
Files changed (1) hide show
  1. app.py +20 -12
app.py CHANGED
@@ -1,28 +1,36 @@
1
  import gradio as gr
2
  from transformers import AutoTokenizer, AutoModel
 
3
 
4
  def chat(prompt):
5
 
6
  global model, tokenizer
7
- response, history = model.chat(tokenizer,
8
- prompt,
9
- history=None,
10
- max_length=64,
11
- top_p=0.7,
12
- temperature=0.95,
13
- num_beam_groups=1,
14
- num_beams=2
15
- )
 
 
 
16
 
17
- return response
 
 
18
 
19
 
20
 
21
 
22
 
23
  if __name__ == '__main__':
24
- tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
25
- model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True).float()
 
 
26
  model.eval()
27
  iface = gr.Interface(fn=chat, inputs="text", outputs="text")
28
  iface.launch()
 
1
  import gradio as gr
2
  from transformers import AutoTokenizer, AutoModel
3
+ from transformers import GPT2Tokenizer,GPT2LMHeadModel
4
 
5
  def chat(prompt):
6
 
7
  global model, tokenizer
8
+ inputs = tokenizer(prompt, return_tensors='pt')
9
+ generation_output = model.generate(**inputs,
10
+ return_dict_in_generate=True,
11
+ output_scores=True,
12
+ max_length=150,
13
+ # max_new_tokens=80,
14
+ do_sample=True,
15
+ top_p=0.6,
16
+ # num_beams=5,
17
+ eos_token_id=50256,
18
+ pad_token_id=0,
19
+ num_return_sequences=1)
20
 
21
+ for idx, sentence in enumerate(generation_output.sequences):
22
+
23
+ return tokenizer.decode(sentence).split('<|endoftext|>')[0]
24
 
25
 
26
 
27
 
28
 
29
  if __name__ == '__main__':
30
+ hf_model_path = 'IDEA-CCNL/Yuyuan-GPT2-110M-SciFi-Chinese'
31
+ tokenizer = GPT2Tokenizer.from_pretrained(hf_model_path)
32
+ model = GPT2LMHeadModel.from_pretrained(hf_model_path)
33
+
34
  model.eval()
35
  iface = gr.Interface(fn=chat, inputs="text", outputs="text")
36
  iface.launch()