Update main.py

main.py

@@ -1,46 +1,51 @@
+from fastapi.staticfiles import StaticFiles
+from fastapi.responses import FileResponse
+from pydantic import BaseModel
+from fastapi import FastAPI
+
 import os
-from flask import Flask, request, jsonify, render_template
+
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
 import torch
 
-app = Flask("Response API")
-name = "microsoft/DialoGPT-medium"
-# microsoft/DialoGPT-small
-# microsoft/DialoGPT-medium
-# microsoft/DialoGPT-large
+app = FastAPI()
 
-# Load the Hugging Face GPT-2 model and tokenizer
-model = GPT2LMHeadModel.from_pretrained(name)
-tokenizer = GPT2Tokenizer.from_pretrained(name)
+class req(BaseModel):
+  prompt: str
+  length: int
 
-# Using CUDA for an optimal experience
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-model = model.to(device)
+@app.get("/")
+def read_root():
+  return FileResponse(path="templates/index.html", media_type="text/html")
 
-# Open a thing for the API
 @app.post("/api")
-def receive_data():
-  data = request.get_json()
+def read_root(data: req):
+  name = "microsoft/DialoGPT-medium"
+  # microsoft/DialoGPT-small
+  # microsoft/DialoGPT-medium
+  # microsoft/DialoGPT-large
   
-  print("Prompt:", data["prompt"])
-  print("Length:", data["length"])
+  # Load the Hugging Face GPT-2 model and tokenizer
+  model = GPT2LMHeadModel.from_pretrained(name)
+  tokenizer = GPT2Tokenizer.from_pretrained(name)
+
+  print("Prompt:", data.prompt)
+  print("Length:", data.length)
+
+  # Using CUDA for an optimal experience
+  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+  model = model.to(device)
   
-  input_text = data["prompt"]
+  input_text = data.prompt
   
   # Tokenize the input text
   input_ids = tokenizer.encode(input_text, return_tensors="pt")
   
   # Generate output using the model
-  output_ids = model.generate(input_ids, max_length=data["length"], num_beams=5, no_repeat_ngram_size=2)
+  output_ids = model.generate(input_ids, max_length=length, num_beams=5, no_repeat_ngram_size=2)
   generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
   
   answer_data = { "answer": generated_text }
   print("Answered with:", answer_data)
-  return jsonify(answer_data)
-
-# Incase a normal browser opens the page
-@app.get("/")
-def not_api():
-  return render_template("index.html")
-
-app.run(debug=False, port=7860, load_dotenv=True)
+  
+  return answer_data