Update main.py

main.py

@@ -1,8 +1,4 @@
-from fastapi.staticfiles import StaticFiles
-from fastapi.responses import FileResponse
-from pydantic import BaseModel
-from fastapi import FastAPI
-
+from flask import Flask, request, render_template, jsonify
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
 model_name = "facebook/blenderbot-1B-distill"
@@ -16,22 +12,24 @@ model_name = "facebook/blenderbot-1B-distill"
 
 # https://www.youtube.com/watch?v=irjYqV6EebU
 
-app = FastAPI()
+app = Flask("AI API")
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-class req(BaseModel):
-  prompt: str
-
 @app.get("/")
 def read_root():
-  return FileResponse(path="templates/index.html", media_type="text/html")
+  return render_template("index.html")
 
-@app.post("/api")
-def read_root(data: req):
-  print("Prompt:", data.prompt)
+@app.route("/test")
+def test_route():
+  return "This is a test route."
 
-  input_text = data.prompt
+@app.route("/api", methods=["POST"])
+def receive_data():
+  data = request.get_json()
+  print("Prompt:", data["prompt"])
+
+  input_text = data["prompt"]
   
   # Tokenize the input text
   input_ids = tokenizer.encode(input_text, return_tensors="pt")
@@ -41,6 +39,8 @@ def read_root(data: req):
   generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
   
   answer_data = { "answer": generated_text }
-  print("Answer:", generated_text)
+  print("Response:", generated_text)
   
-  return answer_data
+  return jsonify(answer_data)
+
+app.run(host="0.0.0.0", port=25428, debug=False)