Update app.py

app.py

@@ -1,26 +1,20 @@
-from flask import Flask, request, render_template
-from transformers import AutoModelForSequenceClassification, AutoTokenizer
-import torch
+from transformers import pipeline
+from flask import Flask, request, jsonify
 
 app = Flask(__name__)
 
-# Load the pre-trained model and tokenizer
-model_name = "https://huggingface.co/atreyuNew/medical_question_answering_chat_Llama2"  # Change this to your desired model
-model = AutoModelForSequenceClassification.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+# Load the model from Hugging Face Spaces
+qa_pipeline = pipeline(model="atreyuNew/medical_question_answering_chat_Llama2")
 
-@app.route('/')
-def home():
-    return render_template('index.html')
+@app.route("/ask", methods=["POST"])
+def ask_question():
+    data = request.get_json()
+    question = data.get("question")
 
-@app.route('/predict', methods=['POST'])
-def predict():
-    text = request.form['text']
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
-    outputs = model(**inputs)
-    probabilities = torch.softmax(outputs.logits, dim=1).tolist()[0]
+    # Use the model to answer the question
+    answer = qa_pipeline(question)
 
-    return render_template('index.html', text=text, probabilities=probabilities)
+    return jsonify({"answer": answer[0]["answer"]})
 
-if __name__ == '__main__':
-    app.run()
+if __name__ == "__main__":
+    app.run(host="0.0.0.0", port=8080)