File size: 870 Bytes
fcce6a8
699e6ab
 
 
ff88e84
fcce6a8
 
 
699e6ab
 
 
 
 
 
 
 
 
fcce6a8
 
 
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
from fastapi import FastAPI
import numpy as np
from sentence_transformers import CrossEncoder
from typing import List
from pydantic import BaseModel

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)#.tolist()
    # (-scores).argsort().tolist()
    return {"results":scores.tolist()}
    
    # keywords = model.encode(inp.keywords)
    # contents = model.encode(inp.contents)
    # return {"results":np.linalg.norm(contents-keywords).tolist()}