| 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, |
| |
| do_sample=True, |
| top_p=0.6, |
| |
| 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() |
