Update app.py

app.py

@@ -10,28 +10,27 @@ MODEL_NAME = "Alaaeldin/pubmedBERT-demo"
 
 @st.cache_resource
 def load_pipeline():
-    return pipeline("question-answering", model=MODEL_NAME, tokenizer=MODEL_NAME, use_auth_token=access_token)
+    # Use the pipeline for question-answering
+    return pipeline("text-generation", model=MODEL_NAME, tokenizer=MODEL_NAME, use_auth_token=access_token)
 
 # Load the pipeline
 qa_pipeline = load_pipeline()
 
 # Streamlit app UI
 st.title("PubMed BERT Q&A App")
-st.write("Ask questions based on biomedical content!")
+st.write("Ask questions directly based on the model's training!")
 
-# User inputs
-context = st.text_area("Enter the biomedical context (e.g., PubMed abstract):", height=200)
+# User input for the question
 question = st.text_input("Enter your question:")
 
 # Button to get the answer
 if st.button("Get Answer"):
-    if context.strip() and question.strip():
-        with st.spinner("Finding the answer..."):
-            result = qa_pipeline(question=question, context=context)
-            st.success(f"Answer: {result['answer']}")
-            st.write(f"Confidence: {result['score']:.2f}")
+    if question.strip():
+        with st.spinner("Generating the answer..."):
+            result = qa_pipeline(question, max_length=100, num_return_sequences=1)
+            st.success(f"Answer: {result[0]['generated_text']}")
     else:
-        st.warning("Please provide both context and a question.")
+        st.warning("Please enter a question.")
 
 # Footer
 st.markdown("---")