doomslayer1434/aimedical_chatbot

1from flask import Flask, request, jsonify
2from flask_cors import CORS  # Import CORS
3import torch
4from transformers import T5Tokenizer, T5ForConditionalGeneration
5
6app = Flask(__name__)
7CORS(app)  # Enable CORS for all routes
8
9# Load model and tokenizer
10model_path = "./t5_chatbot_model"
11tokenizer_path = "./t5_chatbot_tokenizer"
12tokenizer = T5Tokenizer.from_pretrained(tokenizer_path)
13model = T5ForConditionalGeneration.from_pretrained(model_path)
14model.eval()
15
16# Function to generate response
17def generate_response(question, max_length=64, top_k=50, top_p=0.95, temperature=1.0):
18    formatted_question = f"Answer the following question: {question}"
19    inputs = tokenizer(formatted_question, return_tensors="pt", padding=True, truncation=True, max_length=128)
20    outputs = model.generate(
21        input_ids=inputs["input_ids"],
22        attention_mask=inputs["attention_mask"],
23        max_length=max_length,
24        do_sample=True,
25        top_k=top_k,
26        top_p=top_p,
27        temperature=temperature,
28        pad_token_id=tokenizer.pad_token_id,
29    )
30    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
31    return response
32
33@app.route("/chat", methods=["POST", "OPTIONS"])
34def chat():
35    if request.method == "OPTIONS":
36        return jsonify({"message": "CORS preflight request success"}), 200
37
38    try:
39        data = request.get_json()
40        question = data.get("question")
41        if not question:
42            return jsonify({"error": "Missing 'question' in request body"}), 400
43        
44        response = generate_response(question)
45        return jsonify({"response": response})
46    except Exception as e:
47        return jsonify({"error": str(e)}), 500
48
49if __name__ == "__main__":
50    app.run(host="0.0.0.0", port=5000, debug=True)
51