from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, root_validator
from transformers import AutoModel
from typing import List
import os, platform , time



model = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-en',  trust_remote_code=True)

app = FastAPI()

class Validation(BaseModel):
    prompt: List[str]


#Endpoint
@app.post("/jina_embedding")
async def generate_embeddings(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)
    }