Update TextGen/router.py

TextGen/router.py

@@ -1,6 +1,26 @@
 from pydantic import BaseModel
 from fastapi.middleware.cors import CORSMiddleware
+from transformers import AutoModel
+from typing import List
+import os, platform, time
 
+# Your model loading and app initialization code
+model = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-en', trust_remote_code=True)
+
+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
 
@@ -15,3 +35,8 @@ app.add_middleware(
 @app.get("/", tags=["Home"])
 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)