Jackss commited on
Commit
0810777
1 Parent(s): f4e792a
Files changed (1) hide show
  1. main.py +1 -1
main.py CHANGED
@@ -32,7 +32,7 @@ def similarity(input: Input):
32
  inputs = tokenizer(title_abs, padding=True, truncation=True, return_tensors="pt", max_length=512)
33
  result = model(**inputs)
34
  # take the first token in the batch as the embedding
35
- embeddings = result.last_hidden_state[:, 0, :]
36
  res = cosine_similarity(embeddings, embeddings).tolist()
37
  return {"output": res}
38
 
 
32
  inputs = tokenizer(title_abs, padding=True, truncation=True, return_tensors="pt", max_length=512)
33
  result = model(**inputs)
34
  # take the first token in the batch as the embedding
35
+ embeddings = result.last_hidden_state[:, 0, :].detach().numpy()
36
  res = cosine_similarity(embeddings, embeddings).tolist()
37
  return {"output": res}
38