TextGen/router.py

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]

from TextGen import app

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@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(item.prompt)
    }