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

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

@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)