Update TextGen/router.py

TextGen/router.py

@@ -10,18 +10,6 @@ model = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-en', trust_rem
 class Validation(BaseModel):
     prompt: List[str]
 
-def generate_embeddings(x):
-    start_time = time.time()
-    embeddings = model.encode(x).tolist()
-    end_time = time.time()
-    time_taken = end_time - start_time  # Calculate the time taken
-
-    return {
-        "embeddings": embeddings,
-        "time_taken": f"{time_taken:.2f} seconds",
-        "Number_of_sentence_processed": len(x)
-    }
-
 from TextGen import app
 
 app.add_middleware(
@@ -37,6 +25,15 @@ def api_home():
     return {'detail': 'Welcome to FastAPI TextGen Tutorial!'}
 
 
-@app.post("/api/generate", summary="Generate embeddings", tags=["Generate"], response_model=Validation)
-def inference(input_prompt: List[str]):
-    return generate_embeddings(prompt=input_prompt)
+@app.post("/api/generate", summary="Generate embeddings", tags=["Generate"])
+def inference(item: Validation):
+    start_time = time.time()
+    embeddings = model.encode(item.prompt).tolist()
+    end_time = time.time()
+    time_taken = end_time - start_time  # Calculate the time taken
+
+    return {
+        "embeddings": embeddings,
+        "time_taken": f"{time_taken:.2f} seconds",
+        "Number_of_sentence_processed": len(x)
+    }