Loading model once and error handling

app.py

@@ -3,8 +3,6 @@ import torch
 from huggingface_hub import InferenceClient
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-from langchain_core.prompts import PromptTemplate
-from langchain_core.output_parsers import StrOutputParser
 
 # Load HF_TOKEN securely 
 hf_token = os.getenv("HF_TOKEN") 
@@ -16,38 +14,57 @@ client = InferenceClient(api_key=hf_token)
 model_id = "mistralai/Mistral-7B-Instruct-v0.3"
 bart_model_path = "ChijoTheDatascientist/summarization-model"
 
-# Load BART model for summarization
-device = torch.device('cpu')
-bart_tokenizer = AutoTokenizer.from_pretrained(bart_model_path)
-bart_model = AutoModelForSeq2SeqLM.from_pretrained(bart_model_path).to(device)
+# Cache the BART model and tokenizer
+@st.cache_resource
+def load_summarization_model():
+    device = torch.device('cpu')  
+    tokenizer = AutoTokenizer.from_pretrained(bart_model_path)
+    model = AutoModelForSeq2SeqLM.from_pretrained(bart_model_path).to(device)
+    return tokenizer, model
 
+# Load the model and tokenizer 
+bart_tokenizer, bart_model = load_summarization_model()
+
+# Summarize reviews
 @st.cache_data
 def summarize_review(review_text):
-    inputs = bart_tokenizer(review_text, max_length=1024, truncation=True, return_tensors="pt")
-    summary_ids = bart_model.generate(inputs["input_ids"], max_length=40, min_length=10, length_penalty=2.0, num_beams=8, early_stopping=True)
-    summary = bart_tokenizer.decode(summary_ids[0], skip_special_tokens=True)
-    return summary
+    try:
+        if len(review_text) > 1000:
+            return "The review is too long for summarization. Please limit your text to about 1,000 characters, thank you!."
+        inputs = bart_tokenizer(review_text, max_length=1024, truncation=True, return_tensors="pt")
+        summary_ids = bart_model.generate(
+            inputs["input_ids"], 
+            max_length=40, 
+            min_length=10, 
+            length_penalty=2.0, 
+            num_beams=8, 
+            early_stopping=True
+        )
+        summary = bart_tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+        return f"Your review has been successfully summarized! Check the result below:\n\n{summary}"
+    except Exception as e:
+        return f"Something went wrong during the summarization process. Please try again. Error: {e}"
+
 
+# Generate response 
 def generate_response(system_message, user_input, chat_history, max_new_tokens=128):
     try:
         # Prepare the messages for the Hugging Face Inference API
         messages = [{"role": "user", "content": user_input}]
-        
-        # Call the Inference API
         completion = client.chat.completions.create(
             model=model_id, 
             messages=messages, 
             max_tokens=max_new_tokens,
         )
-
-        # Get the response from the API
         response = completion.choices[0].message["content"]
         return response
-
+    except ConnectionError:
+        return "we're having trouble connecting to the server. Please try again later."
     except Exception as e:
-        return f"Error generating response: {e}"
+        return f"Oops! Something went wrong: {e}"
+
 
-# Streamlit app configuration
+# App configuration
 st.set_page_config(page_title="Insight Snap & Summarizer")
 st.title("Insight Snap & Summarizer")
 
@@ -66,34 +83,32 @@ user_input = st.text_area("Enter customer reviews or a question:")
 
 if st.button("Submit"):
     if user_input:
-        # Summarize if the query is feedback-related
-        if "summarize" in user_input.lower():
-            summary = summarize_review(user_input)
-            st.markdown(f"**Summary:** \n{summary}")
-        elif "insight" in user_input.lower() or "feedback" in user_input.lower():
-            system_message = (
-                "You are a helpful assistant providing actionable insights "
-                "from customer feedback to help businesses improve their services."
-            )
-            # Use the last summarized text if available
-            last_summary = st.session_state.get("last_summary", "")
-            query_input = last_summary if last_summary else user_input
-            response = generate_response(system_message, query_input, st.session_state.chat_history)
-            
-            if response:
-                # Update chat history
-                st.session_state.chat_history.append({"role": "user", "content": user_input})
-                st.session_state.chat_history.append({"role": "assistant", "content": response})
-                st.markdown(f"**Insight:** \n{response}")
+        # Show a loading spinner while processing
+        with st.spinner("Processing..."):
+            # Summarize if the query is feedback-related
+            if "summarize" in user_input.lower():
+                summary = summarize_review(user_input)
+                st.markdown(f"**Summary:** \n{summary}")
+            elif "insight" in user_input.lower() or "feedback" in user_input.lower():
+                system_message = (
+                    "You are a helpful assistant providing actionable insights "
+                    "from customer feedback to help businesses improve their services."
+                )
+                last_summary = st.session_state.get("last_summary", "")
+                query_input = last_summary if last_summary else user_input
+                response = generate_response(system_message, query_input, st.session_state.chat_history)
+
+                if response:
+                    st.session_state.chat_history.append({"role": "user", "content": user_input})
+                    st.session_state.chat_history.append({"role": "assistant", "content": response})
+                    st.markdown(f"**Insight:** \n{response}")
+                else:
+                    st.warning("No response generated. Please try again later.")
             else:
-                st.warning("No response generated. Please try again later.")
-        else:
-            st.warning("Please specify if you want to 'summarize' or get 'insights'.")
+                st.warning("Please specify if you want to 'summarize' or get 'insights'.")
 
-        # Store the last summary for insights
-        if "summarize" in user_input.lower():
-            st.session_state["last_summary"] = summary
+           
+            if "summarize" in user_input.lower():
+                st.session_state["last_summary"] = summary
     else:
         st.warning("Please enter customer reviews or ask for insights.")
-
-