File size: 1,277 Bytes
fcce6a8 699e6ab ff88e84 fcce6a8 7c49535 699e6ab 7c49535 699e6ab 197e027 699e6ab fcce6a8 699e6ab 7c49535 2a5588e 699e6ab 7c49535 157594b 7c49535 2a5588e 699e6ab |
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 |
from fastapi import FastAPI
import numpy as np
from sentence_transformers import CrossEncoder
from typing import List
from pydantic import BaseModel
app = FastAPI()
class InputListModel(BaseModel):
keywords: List[str]
contents: List[str]
class InputModel(BaseModel):
keyword: str
content: str
model = CrossEncoder(
# "jinaai/jina-reranker-v2-base-multilingual",
"Alibaba-NLP/gte-multilingual-reranker-base",
trust_remote_code=True,
)
@app.get("/")
def greet_json():
return {"Hello": "World!"}
@app.post("/predict_list")
async def predict_list(inp : InputListModel):
sentence_pairs = [[query, doc] for query,doc in zip(inp.keywords, inp.contents)]
scores = model.predict(sentence_pairs, convert_to_tensor=False)#.tolist()
# (-scores).argsort().tolist()
return {"results":scores.tolist()}
@app.post("/predict")
async def predict(inp : InputModel):
sentence_pairs = [[inp.keyword, inp.content]]
scores = model.predict(sentence_pairs, convert_to_tensor=False)#.tolist()
# (-scores).argsort().tolist()
return {"results":scores.tolist()[0]}
# keywords = model.encode(inp.keywords)
# contents = model.encode(inp.contents)
# return {"results":np.linalg.norm(contents-keywords).tolist()}
|