Update app.py

app.py

@@ -6,10 +6,15 @@ from pydantic import BaseModel
 
 app = FastAPI()
 
-class InputModel(BaseModel):
+class InputListModel(BaseModel):
     keywords: List[str]
     contents: List[str]
 
+class InputModel(BaseModel):
+    keyword: str
+    content: str
+
+
 model = CrossEncoder(
     "jinaai/jina-reranker-v2-base-multilingual",    
     trust_remote_code=True,
@@ -19,12 +24,20 @@ model = CrossEncoder(
 def greet_json():
     return {"Hello": "World!"}
 
-@app.post("/predict")
-async def predict(inp : InputModel):
+@app.post("/predict_list")
+async def predict(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_list")
+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()}
     
     # keywords = model.encode(inp.keywords)
     # contents = model.encode(inp.contents)