switch to fasapi

main.py

@@ -1,22 +1,28 @@
 
-from flask import Flask
+from fastai import FastApi
+from fastapi.middleware.cors import CORSMiddleware
 from generator import generation_function
-from flask import request
-from flask_cors import CORS
-import asyncio
+from pydantic import BaseModel
 
-app = Flask(__name__)
-CORS(app)
 
+app = FastApi()
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=['*'],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
-@app.route('/')
-def hello():
-    return 'Hello world!'
 
+class Prompt(BaseModel):
+    prompts: str
+    lang: str
 
-@app.route('/generate', methods=['POST'])
-def generate():
-    request_data = request.get_json()
+
+@app.post('/generate')
+def generate(prompt: Prompt):
+    request_data = prompt.dict()
 
     text = request_data["prompts"]
 
@@ -27,7 +33,3 @@ def generate():
         generated_recepie = generation_function(text)
 
     return generated_recepie
-
-
-if __name__ == '__main__':
-    app.run(host="0.0.0.0", port=7860)