Update main.py

main.py

@@ -1,7 +1,10 @@
-from flask import Flask, request, jsonify
 from torch import Tensor
 from transformers import AutoTokenizer, AutoModel
 from ctranslate2 import Translator
+from typing import Union
+
+from fastapi import FastAPI
+from pydantic import BaseModel
 
 
 def average_pool(last_hidden_states: Tensor,
@@ -13,24 +16,40 @@ def average_pool(last_hidden_states: Tensor,
 
 # text-ada replacement
 embeddingTokenizer = AutoTokenizer.from_pretrained(
-    './multilingual-e5-base')
-embeddingModel = AutoModel.from_pretrained('./multilingual-e5-base')
+    './models/multilingual-e5-base')
+embeddingModel = AutoModel.from_pretrained('./models/multilingual-e5-base')
 
 # chatGpt replacement
 inferenceTokenizer = AutoTokenizer.from_pretrained(
-    "./ct2fast-flan-alpaca-xl")
+    "./models/ct2fast-flan-alpaca-xl")
 inferenceTranslator = Translator(
-    "./ct2fast-flan-alpaca-xl", compute_type="int8", device="cpu")
+    "./models/ct2fast-flan-alpaca-xl", compute_type="int8", device="cpu")
+
+
+class EmbeddingRequest(BaseModel):
+    input: Union[str, None] = None
+
+
+class TokensCountRequest(BaseModel):
+    input: Union[str, None] = None
+
+
+class InferenceRequest(BaseModel):
+    input: Union[str, None] = None
+    max_length: Union[int, None] = 0
+
 
+app = FastAPI()
 
-app = Flask(__name__)
 
+@app.get("/")
+async def root():
+    return {"message": "Hello World"}
 
-@app.route('/text-embedding', methods=['POST'])
-def text_embedding():
-    # Get the JSON data from the request
-    data = request.get_json()
-    input = data["input"]
+
+@app.post("/text-embedding")
+async def text_embedding(request: EmbeddingRequest):
+    input = request.input
 
     # Process the input data
     batch_dict = embeddingTokenizer([input], max_length=512,
@@ -38,28 +57,24 @@ def text_embedding():
     outputs = embeddingModel(**batch_dict)
     embeddings = average_pool(outputs.last_hidden_state,
                               batch_dict['attention_mask'])
-    token_ids = batch_dict["input_ids"][0].tolist()
 
-    # Create a JSON response
-    response = {
+    # create response
+    return {
         'embedding': embeddings[0].tolist()
     }
 
-    return jsonify(response)
 
-
-@app.route('/inference', methods=['POST'])
-def inference():
-    # Get the JSON data from the request
-    data = request.get_json()
-    input_text = data["input"]
+@app.post('/inference')
+async def inference(request: InferenceRequest):
+    input_text = request.input
     max_length = 256
     try:
-        max_length = int(data["max_length"])
+        max_length = int(request.max_length)
         max_length = min(1024, max_length)
     except:
         pass
 
+    # process request
     input_tokens = inferenceTokenizer.convert_ids_to_tokens(
         inferenceTokenizer.encode(input_text))
 
@@ -70,31 +85,21 @@ def inference():
     output_text = inferenceTokenizer.decode(
         inferenceTokenizer.convert_tokens_to_ids(output_tokens))
 
-    # Create a JSON response
-    response = {
+    # create response
+    return {
         'generated_text': output_text
     }
 
-    return jsonify(response)
 
-
-@app.route('/tokens-count', methods=['POST'])
-def tokens_count():
-    # Get the JSON data from the request
-    data = request.get_json()
-    input_text = data["input"]
+@app.post('/tokens-count')
+async def tokens_count(request: TokensCountRequest):
+    input_text = request.input
 
     tokens = inferenceTokenizer.convert_ids_to_tokens(
         inferenceTokenizer.encode(input_text))
 
-    # Create a JSON response
+    # create response
     response = {
         'tokens': tokens,
         'total': len(tokens)
     }
-
-    return jsonify(response)
-
-
-if __name__ == '__main__':
-    app.run()