Update app.py

app.py

@@ -1,27 +1,17 @@
 import streamlit as st
-import logging
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
-# Set the logger to display only CRITICAL messages
-logging.basicConfig(level=logging.CRITICAL)
-
-
-# Cache the model and tokenizer to avoid reloading it every time
-
-def load_model():
+# Streamlit's cache decorator to cache the model and tokenizer loading
+@st.cache(allow_output_mutation=True)
+def load_pipeline():
     model_name = "NousResearch/Llama-2-7b-chat-hf"  # Replace with your actual model name
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForCausalLM.from_pretrained(model_name)
-    return model, tokenizer
-
-model, tokenizer = load_model()
+    chat_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer, max_length=300)
+    return chat_pipeline
 
-# Function to generate text with the model
-def generate_text(prompt):
-    formatted_prompt = f"[INST] {prompt} [/INST]"  # Format the prompt according to your specification
-    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_length=300)
-    result = pipe(formatted_prompt)
-    return result[0]['generated_text']
+# Initialize the pipeline
+chat_pipeline = load_pipeline()
 
 st.title("Interact with Your Model")
 
@@ -30,8 +20,11 @@ user_input = st.text_area("Enter your prompt:", "")
 
 if st.button("Submit"):
     if user_input:
-        # Generate text based on the input
-        generated_text = generate_text(user_input)
-        st.write(generated_text)
+        try:
+            # Generate text based on the input
+            generated_text = chat_pipeline(user_input)[0]['generated_text']
+            st.write(generated_text)
+        except Exception as e:
+            st.error(f"Error generating text: {e}")
     else:
-        st.write("Please enter a prompt.")
+        st.write("Please enter a prompt.")