File size: 739 Bytes
7430631 8d520f8 7430631 8d520f8 be4dd45 7430631 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 |
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}
|