Update app.py

app.py

@@ -5,49 +5,38 @@ from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 st.title('Chest and Physical Limitations LLM Query')
 
 # Initialize variables
-model_name = "Abbeite/chest_and_physical_limitations2"
+model_name = "Abbeite/chest_and_physical_limitations"
+generator = None
 
-# Use st.cache to cache the model loading function
-@st.cache(allow_output_mutation=True, show_spinner=True)
+# Function to check model existence and load it
 def load_model(model_name):
     try:
         tokenizer = AutoTokenizer.from_pretrained(model_name)
         model = AutoModelForCausalLM.from_pretrained(model_name)
         generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
-        return generator, None  # Return None for the error message on success
+        return generator
     except Exception as e:
-        return None, str(e)  # Return None for the generator and the exception message on failure
-    return generator, None  # Return the generator and None for the error
-
-# Load the model and check for errors
-# Attempt to load the model and receive an error message if loading fails
-generator, error_message = load_model(model_name)
-
-# If an error message was returned, display it using st.error
-if error_message:
-    st.error(f"Failed to load model {model_name}: {error_message}")
+        st.error(f"Failed to load model {model_name}: {str(e)}")
+        return None
+
+# Load the model
+generator = load_model(model_name)
+
+# User prompt input
+if generator:  # Proceed only if the model is successfully loaded
+    user_prompt = st.text_area("Enter your prompt here:")
+
+    # Button to generate text
+    if st.button('Generate'):
+        if user_prompt:
+            # Generate response
+            try:
+                response = generator(user_prompt, max_length=50, clean_up_tokenization_spaces=True)
+                # Display the generated text
+                st.text_area("Response:", value=response[0]['generated_text'], height=250, disabled=True)
+            except Exception as e:
+                st.error(f"Error generating response: {str(e)}")
+        else:
+            st.warning("Please enter a prompt.")
 else:
-    # Proceed with the rest of your Streamlit app logic if the model was loaded successfully
-    user_input_section(generator)
-
-# User prompt input section wrapped in a function to control reruns
-def user_input_section(generator):
-    if generator:  # Check moved inside the function
-        user_prompt = st.text_area("Enter your prompt here:")
-
-        # Button to generate text
-        if st.button('Generate'):
-            if user_prompt:
-                # Generate response
-                try:
-                    response = generator(user_prompt, max_length=100, clean_up_tokenization_spaces=True)
-                    # Display the generated text
-                    st.text_area("Response:", value=response[0]['generated_text'], height=250, disabled=True)
-                except Exception as e:
-                    st.error(f"Error generating response: {str(e)}")
-            else:
-                st.warning("Please enter a prompt.")
-    else:
-        st.error("Model could not be loaded. Please ensure the model name is correct and try again.")
-
-user_input_section(generator)
+    st.error("Model could not be loaded. Please ensure the model name is correct and try again.")