ChatFinance / app.py
Joe99's picture
first finance bot
7a7a56a
raw
history blame
No virus
1.51 kB
import transformers
import gradio as gr
import warnings
import torch
warnings.simplefilter('ignore')
device = "cuda" if torch.cuda.is_available() else "cpu"
tokenizer = transformers.GPT2Tokenizer.from_pretrained('gpt2')
#add padding token, beginstring and endstring tokens
tokenizer.add_special_tokens(
{
"pad_token":"<pad>",
"bos_token":"<startstring>",
"eos_token":"<endstring>"
})
#add bot token since it is not a special token
tokenizer.add_tokens(["<bot>:"])
model = transformers.GPT2LMHeadModel.from_pretrained('gpt2')
model.resize_token_embeddings(len(tokenizer))
model.load_state_dict(torch.load('gpt2talk.pt', map_location=torch.device('cpu')))
model.eval()
def inference(quiz):
quiz1 = quiz
quiz = "<startstring>"+quiz+" <bot>:"
quiztoken = tokenizer(quiz,
return_tensors='pt'
)
answer = model.generate(**quiztoken, max_length=200, top_k=0.7,top_p=0.1)[0]
answer = tokenizer.decode(answer, skip_special_tokens=True)
answer = answer.replace(" <bot>:","").replace(quiz1,"") + '.'
return answer
def chatbot(input_text):
response = inference(input_text)
return response
# Create the Gradio interface
iface = gr.Interface(
fn=chatbot,
inputs=gr.Textbox(),
outputs=gr.Textbox(),
live=False, #set false to avoid caching
interpretation="chat",
title="ChatFinance",
description="Ask the a question and see its response!",
)
# Launch the Gradio interface
iface.launch()