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

model = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-en',token = os.environ.get("hf_token"), trust_remote_code=True)

app = FastAPI()

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


#Endpoint
@app.post("/jina_embedding")
async def chaatie_agent(item: Validation):
    # Assuming model.encode returns a list of numpy arrays (one for each prompt)
    embeddings = model.encode(item.prompt)

    # Convert each numpy array in the list to a list
    embeddings_list = [embedding.tolist() for embedding in embeddings]

    return {"embeddings": embeddings_list}