Moved from sentence-transformers to pure transformers

core.py

@@ -1,10 +1,13 @@
 import streamlit as st
-from sentence_transformers import SentenceTransformer
 from huggingface_hub import HfApi, HfFolder
 import datasets
 import logging
 import os
 
+from transformers import AutoTokenizer, AutoModel
+import torch
+import torch.nn.functional as F
+
 
 @st.cache_data
 def login():
@@ -12,7 +15,7 @@ def login():
         logging.info("Trying to log in to HF")
         st.session_state['logged'] = True
         HF_TOKEN = os.environ.get("HF_TOKEN")
-        api=HfApi()
+        api = HfApi()
         api.set_access_token(HF_TOKEN)
         folder = HfFolder()
         folder.save_token(HF_TOKEN)
@@ -26,16 +29,18 @@ def login():
 @st.cache_resource
 def load_model():
     logging.info("Trying to load model")
-    model = SentenceTransformer('sentence-transformers/paraphrase-MiniLM-L6-v2')
+    tokenizer = AutoTokenizer.from_pretrained(
+        'sentence-transformers/all-MiniLM-L6-v2')
+    model = AutoModel.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
     logging.info("Model loaded")
-    return model
+    return model, tokenizer
 
 
 @st.cache_resource
 def load_index():
     logging.info("Trying to load index")
     index = datasets.load_dataset(
-        "eremeev-d/arxiv-abstracts-small", 
+        "eremeev-d/arxiv-abstracts-small",
         use_auth_token=True,
         split="train"
     )
@@ -44,11 +49,29 @@ def load_index():
     return index
 
 
+def mean_pooling(model_output, attention_mask):
+    token_embeddings = model_output[0]
+    input_mask_expanded = attention_mask.unsqueeze(-1) \
+        .expand(token_embeddings.size()).float()
+    return torch.sum(token_embeddings * input_mask_expanded, 1) \
+        / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
+
+
+def get_embedding(query, model, tokenizer):
+    encoded_input = tokenizer(
+        query, padding=True, truncation=True, return_tensors='pt')
+    with torch.no_grad():
+        embeds = model(**encoded_input)
+    embeds = mean_pooling(embeds, encoded_input['attention_mask'])
+    embeds = F.normalize(embeds, p=2, dim=1)
+    return embeds
+
+
 def get_answers(query):
     logging.info("Getting answers for {}".format(query))
-    model = load_model()
+    model, tokenizer = load_model()
     index = load_index()
-    query_embedding = model.encode(query)
+    query_embedding = get_embedding(query, model, tokenizer).reshape(-1).numpy()
     scores, answers = index.get_nearest_examples('embedding', query_embedding)
     logging.info("Succesfully got answers for {}".format(query))
     return answers

requirements.txt

@@ -1,4 +1,5 @@
 faiss-cpu~=1.7.2
 sentence-transformers~=2.2.2
 datasets~=2.10.1
-huggingface_hub~=0.10.1
+huggingface_hub~=0.10.1
+torch