Upload app.py with huggingface_hub

app.py

@@ -1,15 +1,23 @@
 
 import gradio as gr
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
-model_name = "AI4FinanceFoundation/FinGPT-Foundation"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
+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)
 
-def answer(query):
-    response = pipe(query, max_new_tokens=300, do_sample=True)[0]["generated_text"]
-    return response
+chat_history = []
 
-gr.Interface(fn=answer, inputs="text", outputs="text", title="FinGPT QA").launch()
+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()