Update app.py

app.py

@@ -1,26 +1,42 @@
 import streamlit as st
-from transformers import pipeline
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 
 # Title of the web application
 st.title('Chest and Physical Limitations LLM Query')
 
-# Initialize the model pipeline
-# Ensure to use the correct model identifier
+# Initialize variables
 model_name = "Abbeite/chest_and_physical_limitations"
-generator = pipeline('text-generation', model=model_name)
+generator = None
+
+# 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
+    except Exception as e:
+        st.error(f"Failed to load model {model_name}: {str(e)}")
+        return None
+
+# Load the model
+generator = load_model(model_name)
 
 # User prompt input
-user_prompt = st.text_area("Enter your prompt here:")
+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.")
+    # 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:
+    st.error("Model could not be loaded. Please ensure the model name is correct and try again.")