Spaces:
Runtime error
Runtime error
| # 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 | |
| def api_home(): | |
| return {'hello': 'Welcome!'} | |
| # API to 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 | |
| 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)}") | |