Update app.py

app.py

@@ -5,10 +5,14 @@ import io
 from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
 
 
-# Load fine-tuned model and tokenizer (adjust the model name)
-model_name = "TAgroup5/news-classification-model"  # Replace with the correct model name
-model = AutoModelForSequenceClassification.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+# Load fine-tuned models and tokenizers for both functions 
+model_name_classification = "TAgroup5/news-classification-model"  # Replace with the correct model name
+model = AutoModelForSequenceClassification.from_pretrained(model_name_classification)
+tokenizer = AutoTokenizer.from_pretrained(model_name_classification)
+
+model_name_qa = "distilbert-base-cased-distilled-squad"
+model_qa = AutoModelForQuestionAnswering.from_pretrained(model_name_qa)
+tokenizer_qa = AutoTokenizer.from_pretrained(model_name_qa)
 
 # Initialize pipelines
 text_classification_pipeline = pipeline("text-classification", model=model, tokenizer=tokenizer)
@@ -68,8 +72,8 @@ question = st.text_input("Ask a question:")
 context = st.text_area("Provide the news article or content for the Q&A:", height=150)
 
 if question and context.strip():
-    qa_model_name = "distilbert-base-uncased-distilled-squad"  # Example of a common Q&A model
-    qa_pipeline = pipeline("question-answering", model=qa_model_name, tokenizer=qa_model_name)
+    model_name_qa = "distilbert-base-uncased-distilled-squad"  # Example of a common Q&A model
+    qa_pipeline = pipeline("question-answering", model=model_name_qa, tokenizer=model_name_qa)
     result = qa_pipeline(question=question, context=context)
     
     # Check if the result contains an answer