Update app.py

app.py

@@ -5,10 +5,12 @@ import torch
 # Initialize the model and tokenizer
 model_id = "meta-llama/Meta-Llama-3-8B"
 
-@st.cache(allow_output_mutation=True)
+@st.cache_resource
 def load_model():
     tokenizer = AutoTokenizer.from_pretrained(model_id)
     model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16)
+    tokenizer.add_special_tokens({"pad_token": "<pad>"})
+    model.resize_token_embeddings(len(tokenizer))
     return pipeline("text-generation", model=model, tokenizer=tokenizer, device_map="auto")
 
 pipe = load_model()
@@ -22,8 +24,11 @@ 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'])
+            try:
+                result = pipe(user_input, max_length=100, num_return_sequences=1, truncation=True, pad_token_id=tokenizer.pad_token_id)
+                st.success("Here are your insights:")
+                st.write(result[0]['generated_text'])
+            except Exception as e:
+                st.error(f"An error occurred: {e}")
     else:
-        st.error("Please enter a query.")
+        st.error("Please enter a query.")