File size: 1,310 Bytes
4d42d17
e646b3e
f07a852
4d42d17
2e22a78
 
e646b3e
f07a852
 
 
 
 
 
 
 
 
 
 
 
2e22a78
f07a852
 
 
4d42d17
e646b3e
 
 
 
 
f07a852
 
 
 
e646b3e
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
import gradio as gr
from transformers import AutoTokenizer, AutoModel
from transformers import GPT2Tokenizer,GPT2LMHeadModel

def chat(prompt):

    global model, tokenizer
    inputs = tokenizer(prompt, return_tensors='pt')
    generation_output = model.generate(**inputs,
                                       return_dict_in_generate=True,
                                       output_scores=True,
                                       max_length=150,
                                       # max_new_tokens=80,
                                       do_sample=True,
                                       top_p=0.6,
                                       # num_beams=5,
                                       eos_token_id=50256,
                                       pad_token_id=0,
                                       num_return_sequences=1)

    for idx, sentence in enumerate(generation_output.sequences):

              return tokenizer.decode(sentence).split('<|endoftext|>')[0]





if __name__ == '__main__':
    hf_model_path = 'IDEA-CCNL/Yuyuan-GPT2-110M-SciFi-Chinese'
    tokenizer = GPT2Tokenizer.from_pretrained(hf_model_path)
    model = GPT2LMHeadModel.from_pretrained(hf_model_path)

    model.eval()
    iface = gr.Interface(fn=chat, inputs="text", outputs="text")
    iface.launch()