import gradio as gr | |
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline | |
tokenizer = AutoTokenizer.from_pretrained("FinGPT/fingpt-mt_llama2-7b_lora", trust_remote_code=True) | |
model = AutoModelForCausalLM.from_pretrained("FinGPT/fingpt-mt_llama2-7b_lora", trust_remote_code=True, device_map="auto") | |
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer) | |
chat_history = [] | |
def chat(user_input, history): | |
prompt = (history + "\nUser: " + user_input) if history else ("User: " + user_input) | |
output = pipe(prompt, max_new_tokens=300, do_sample=True)[0]["generated_text"] | |
history = prompt + "\nAssistant: " + output | |
return output, history | |
with gr.Blocks() as demo: | |
chatbot = gr.Chatbot() | |
txt = gr.Textbox(placeholder="Ask a finance question...", show_label=False) | |
state = gr.State("") | |
txt.submit(lambda msg, hist: (chatbot + [(msg, chat(msg, hist)[0])], chat(msg, hist)[1]), [txt, state], [chatbot, state]) | |
demo.launch() | |