rag-demo / main.py
moraxgiga's picture
Update main.py
8d520f8 verified
raw
history blame
739 Bytes
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}