|
from fastapi import FastAPI |
|
import numpy as np |
|
from sentence_transformers import CrossEncoder |
|
from typing import List |
|
|
|
app = FastAPI() |
|
|
|
class InputModel(BaseModel): |
|
keywords: List[str] |
|
contents: List[str] |
|
|
|
model = CrossEncoder( |
|
"jinaai/jina-reranker-v2-base-multilingual", |
|
trust_remote_code=True, |
|
) |
|
|
|
@app.get("/") |
|
def greet_json(): |
|
return {"Hello": "World!"} |
|
|
|
@app.post("/predict") |
|
def predict(inp : InputModel): |
|
sentence_pairs = [[query, doc] for query,doc in zip(inp.keywords, inp.contents)] |
|
scores = model.predict(sentence_pairs, convert_to_tensor=False) |
|
|
|
return {"results":scores.tolist()} |
|
|
|
|
|
|
|
|
|
|