Update app.py

app.py

@@ -1,28 +1,29 @@
-import os
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import streamlit as st
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 import torch
 
-# Set the Hugging Face token
-hf_token = os.getenv("HF_TOKEN")
+# Initialize the model and tokenizer
+model_id = "meta-llama/Meta-Llama-3-8B"
 
-# Load the model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B", use_auth_token=hf_token)
-model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B", use_auth_token=hf_token, torch_dtype=torch.bfloat16)
+@st.cache(allow_output_mutation=True)
+def load_model():
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16)
+    return pipeline("text-generation", model=model, tokenizer=tokenizer, device_map="auto")
 
-st.title("FinWise AI")
-st.write("Enter a message to generate a response:")
+pipe = load_model()
 
-messages = st.text_area("Message", "Hey, how are you doing today?")
-if st.button("Generate Response"):
-    with st.spinner("Generating response..."):
-        inputs = tokenizer(messages, return_tensors="pt")
-        outputs = model.generate(
-            inputs.input_ids,
-            max_new_tokens=256,
-            do_sample=True,
-            temperature=0.6,
-            top_p=0.9
-        )
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        st.write(response)
+# Streamlit app UI
+st.title("FinWise AI 🏆")
+st.write("Your AI-powered financial advisor")
+
+user_input = st.text_area("Enter your query about stock market investments:", "")
+
+if st.button("Get Insights"):
+    if user_input:
+        with st.spinner("Generating insights..."):
+            result = pipe(user_input, max_length=100, num_return_sequences=1)
+            st.success("Here are your insights:")
+            st.write(result[0]['generated_text'])
+    else:
+        st.error("Please enter a query.")