Update app.py

app.py

@@ -13,9 +13,9 @@ st.set_page_config(
 
 # Load Hugging Face models and tokenizer for text generation
 @st.cache_resource
-def load_model(model_name):
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name)
+def load_model():
+    tokenizer = AutoTokenizer.from_pretrained("harishussain12/Disease_Managment")
+    model = AutoModelForCausalLM.from_pretrained("harishussain12/Disease_Managment")
     return tokenizer, model
 
 # Load Hugging Face models and tokenizer for document question answering
@@ -27,8 +27,8 @@ def load_document_model():
 
 # Function to create a text generation pipeline
 @st.cache_resource
-def create_pipeline(model_name):
-    tokenizer, model = load_model(model_name)
+def create_pipeline():
+    tokenizer, model = load_model()
     return pipeline("text-generation", model=model, tokenizer=tokenizer)
 
 # Function to create a document question answering pipeline
@@ -48,6 +48,10 @@ def read_pdf(file):
     except Exception as e:
         return f"Error reading PDF: {e}"
 
+# Load pipelines
+text_pipeline = create_pipeline()
+document_pipeline = create_document_pipeline()
+
 # Custom CSS for styling
 st.markdown(
     """
@@ -145,9 +149,6 @@ with col2:
         with st.spinner("Generating response..."):
             try:
                 if model_selection == "Disease Analysis":
-                    # Adjusted model name to match working configurations
-                    pipe = create_pipeline("harishussain12/Disease_Managment")
-                    
                     context = ""
                     if uploaded_file is not None:
                         file_content = read_pdf(uploaded_file)
@@ -157,14 +158,10 @@ with col2:
                             context = file_content
 
                     query_input = search_input + (f"\n\nContext:\n{context}" if context else "")
-                    st.write(f"Debug: Query Input - {query_input}")  # Debug log for troubleshooting
-
-                    response = pipe(query_input, max_length=200, num_return_sequences=1)
+                    response = text_pipeline(query_input, max_length=200, num_return_sequences=1)
                     st.markdown(f"### Response:\n{response[0]['generated_text']}")
 
                 elif model_selection == "Document Analysis":
-                    pipe = create_document_pipeline()
-
                     context = ""
                     if uploaded_file is not None:
                         file_content = read_pdf(uploaded_file)
@@ -174,7 +171,7 @@ with col2:
                             context = file_content
 
                     if search_input and context:
-                        result = pipe({"question": search_input, "context": context})
+                        result = document_pipeline({"question": search_input, "context": context})
                         st.markdown(f"### Answer:\n{result['answer']}")
 
             except Exception as e: