File size: 984 Bytes
85ea4bd
6b88819
4644a82
85ea4bd
4644a82
 
85ea4bd
6b88819
85ea4bd
4644a82
85ea4bd
4644a82
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

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()