rag-demo / main.py
moraxgiga's picture
Update main.py
be4dd45 verified
raw
history blame
No virus
715 Bytes
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, root_validator
from transformers import AutoModel
from typing import List
model = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-en',token = access_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}