| import gradio as gr | |
| from transformers import GPT2Tokenizer, GPT2LMHeadModel, TrainingArguments, Trainer | |
| import tiktoken | |
| import torch | |
| model_name = "paramasivan27/gpt2_for_q_and_a" | |
| tokenizer = GPT2Tokenizer.from_pretrained(model_name) | |
| model = GPT2LMHeadModel.from_pretrained(model_name) | |
| def ask_question(question): | |
| inputs = tokenizer.encode('Q: ' + question + ' A:', return_tensors='pt') | |
| attention_mask = torch.ones(inputs.shape) | |
| outputs = model.generate(inputs, attention_mask = attention_mask, max_new_tokens=100, num_return_sequences=1) | |
| gen_text = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
| question, answer = gen_text.split(' A:') | |
| return question, answer | |
| """ | |
| For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface | |
| """ | |
| demo = gr.ChatInterface( | |
| ask_question, | |
| additional_inputs=[ | |
| gr.Textbox(value="You are a friendly Chatbot.", label="System message"), | |
| gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"), | |
| gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"), | |
| gr.Slider( | |
| minimum=0.1, | |
| maximum=1.0, | |
| value=0.95, | |
| step=0.05, | |
| label="Top-p (nucleus sampling)", | |
| ), | |
| ], | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |