Stuff

main.py

@@ -32,7 +32,7 @@ def similarity(input: Input):
     inputs = tokenizer(title_abs, padding=True, truncation=True, return_tensors="pt", max_length=512)
     result = model(**inputs)
     # take the first token in the batch as the embedding
-    embeddings = result.last_hidden_state[:, 0, :]
+    embeddings = result.last_hidden_state[:, 0, :].detach().numpy()
     res = cosine_similarity(embeddings, embeddings).tolist()
     return {"output": res}