from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, root_validator
from transformers import AutoModel
from typing import List
import os, platform, time
from fastapi.middleware.cors import CORSMiddleware  # Import CORSMiddleware

# Your model loading and app initialization code
model = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-en', trust_remote_code=True)
app = FastAPI()

# Set up CORS middleware configuration
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Allows all origins
    allow_credentials=True,
    allow_methods=["*"],  # Allows all methods
    allow_headers=["*"],  # Allows all headers
)

# Your Validation model and endpoint definitions remain the same
class Validation(BaseModel):
    prompt: List[str]

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