Stuff

main.py

@@ -4,6 +4,7 @@ from fastapi.responses import FileResponse
 from transformers import AutoTokenizer, AutoModel
 import numpy as np
 from sklearn.metrics.pairwise import cosine_similarity
+from pydantic import BaseModel
 
 # load model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained('allenai/specter')
@@ -15,13 +16,17 @@ model = AutoModel.from_pretrained('allenai/specter')
 # concatenate title and abstract
 
 
+class Input(BaseModel):
+    papers: list = []
+
+
 app = FastAPI()
 
 
 
 @app.post('/similarity')
-def similarity(input):
-    papers = input['papers']
+def similarity(input: Input):
+    papers = input.papers
     title_abs = [d['title'] + tokenizer.sep_token + (d.get('abstract') or '') for d in papers]
     # preprocess the input
     inputs = tokenizer(title_abs, padding=True, truncation=True, return_tensors="pt", max_length=512)