from transformers import pipeline
from flask import Flask, request, jsonify

app = Flask(__name__)

# Load the model from Hugging Face Spaces
qa_pipeline = pipeline(model="atreyuNew/medical_question_answering_chat_Llama2")

@app.route("/ask", methods=["POST"])
def ask_question():
    data = request.get_json()
    question = data.get("question")

    # Use the model to answer the question
    answer = qa_pipeline(question)

    return jsonify({"answer": answer[0]["answer"]})

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=8080)