|
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 |
|
|
|
|
|
model = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-en', trust_remote_code=True) |
|
app = FastAPI() |
|
|
|
|
|
app.add_middleware( |
|
CORSMiddleware, |
|
allow_origins=["*"], |
|
allow_credentials=True, |
|
allow_methods=["*"], |
|
allow_headers=["*"], |
|
) |
|
|
|
|
|
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 |
|
|
|
return { |
|
"embeddings": embeddings, |
|
"time_taken": f"{time_taken:.2f} seconds", |
|
"Number_of_sentence_processed": len(item.prompt) |
|
} |