File size: 1,286 Bytes
2015c1f 9b49b40 2015c1f 9b49b40 2015c1f 9b49b40 2015c1f 9b49b40 2015c1f 9b49b40 2015c1f | 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 | import os
import uvicorn
from fastapi import FastAPI, Request
from sentence_transformers import SentenceTransformer, CrossEncoder
app = FastAPI()
embedding_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
reranker_model = CrossEncoder("BAAI/bge-reranker-large")
@app.get("/")
def root():
return {"message": "Embedding & Reranker API is live!"}
@app.post("/embed")
async def embed(request: Request):
data = await request.json()
text = data.get("text", "")
if not text:
return {"error": "Missing text"}
embedding = embedding_model.encode(text).tolist()
return {"embedding": embedding}
@app.post("/rerank")
async def rerank(request: Request):
data = await request.json()
query = data.get("query")
docs = data.get("documents", [])
if not query or not docs:
return {"error": "Missing query or documents"}
inputs = [[query, doc["content"]] for doc in docs]
scores = reranker_model.predict(inputs)
for doc, score in zip(docs, scores):
doc["rerank_score"] = float(score)
return {"documents": docs}
@app.get("/health")
async def health():
return {"status": "ok"}
if __name__ == "__main__":
port = int(os.environ.get("PORT", 10000))
uvicorn.run(app, host="0.0.0.0", port=port) |