Spaces:
Sleeping
Sleeping
File size: 1,278 Bytes
5f0af37 990c9b3 5f0af37 3f00318 5f0af37 |
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 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 |
# main.py
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List
from vllm import LLM
import numpy as np
# Initialize the model
llm = LLM(model='BAAI/bge-m3', task="embed")
# Initialize FastAPI app
app = FastAPI()
# Define request schemas
class DocumentsRequest(BaseModel):
documents: List[str]
class QueryRequest(BaseModel):
query: str
@app.get("/", tags=["Home"])
def api_home():
return {'hello': 'Welcome!'}
# API to embed documents
@app.post("/embed_documents")
def embed_documents(request: DocumentsRequest):
try:
docs = request.documents
docs_embd = llm.encode(docs)
docs_embd = [doc.outputs.data.numpy().tolist() for doc in docs_embd]
return {"embeddings": docs_embd}
except Exception as e:
raise HTTPException(status_code=500, detail=f"Error embedding documents: {str(e)}")
# API to embed query
@app.post("/embed_query")
def embed_query(request: QueryRequest):
try:
query = request.query
query_embd = llm.encode(query)
query_embd = query_embd[0].outputs.data.numpy().tolist()
return {"embedding": query_embd}
except Exception as e:
raise HTTPException(status_code=500, detail=f"Error embedding query: {str(e)}")
|