Update app.py

app.py

@@ -1,9 +1,10 @@
-from flask import Flask, request, jsonify
+from fastapi import FastAPI
+from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import json
 
-app = Flask(__name__)
+app = FastAPI()
 
 # Load DialoGPT model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-large")
@@ -13,7 +14,10 @@ model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-large")
 with open("uts_courses.json", "r") as file:
     courses_data = json.load(file)
 
-def generate_response(user_input):
+class UserInput(BaseModel):
+    user_input: str
+
+def generate_response(user_input: str):
     if user_input.lower() == "help":
         return "I can help you with information about UTS courses. Feel free to ask!"
     elif user_input.lower() == "exit":
@@ -32,11 +36,7 @@ def generate_response(user_input):
         response = tokenizer.decode(response_ids[0], skip_special_tokens=True)
         return response
 
-@app.route("/", methods=["POST"])
-def chat():
-    user_input = request.json["user_input"]
-    response = generate_response(user_input)
-    return jsonify({"response": response})
-
-if __name__ == "__main__":
-    app.run(debug=True)
+@app.post("/")
+def chat(user_input: UserInput):
+    response = generate_response(user_input.user_input)
+    return {"response": response}