reverting unnecessary changes

app.py

@@ -3,7 +3,7 @@ import datasets
 import faiss
 import os
 
-from transformers import pipeline, AutoModel, AutoTokenizer
+from transformers import pipeline  # , AutoModel, AutoTokenizer
 
 
 auth_token = os.environ.get("CLARIN_KNEXT")
@@ -23,49 +23,6 @@ textbox = gr.Textbox(
 )
 
 
-def prepare_query(tokenizer, query, max_seq_length=300):
-    # temporary solution b/c of padding (which is unnecessary for inference)
-    start_token: str = "[unused0]"
-    end_token: str = "[unused1]"
-    
-    left_context = query.split(start_token)[0]
-    right_context = query.split(end_token)[-1]
-    mention = query.split(start_token)[-1].split(end_token)[0]
-    
-    mention_ids = tokenizer(
-        f"{start_token}{mention}{end_token}",
-        add_special_tokens=False
-    )['input_ids']
-
-    left_ids = tokenizer(left_context, add_special_tokens=False)['input_ids']
-    left_quota = (max_seq_length - len(mention_ids)) // 2 - 1
-
-    right_ids = tokenizer(right_context, add_special_tokens=False)['input_ids']
-    right_quota = max_seq_length - len(mention_ids) - left_quota - 2
-
-    left_add, right_add = len(left_ids), len(right_ids)
-    if left_add <= left_quota:
-        right_quota += left_quota - left_add if right_add > right_quota else 0
-    else:
-        left_quota += right_quota - right_add if right_add <= right_quota else 0
-
-    context_ids = [
-        tokenizer.cls_token_id,
-        *left_ids[-left_quota:],
-        *mention_ids,
-        *right_ids[:right_quota],
-        tokenizer.sep_token_id
-    ]
-
-    padding_length = max_seq_length - len(context_ids)
-    # attention_mask = [1] * len(context_ids) + [0] * padding_length
-
-    context_ids += [tokenizer.pad_token_id] * padding_length
-
-    assert len(context_ids) == max_seq_length
-    return context_ids
-
-
 def load_index(index_data: str = "clarin-knext/entity-linking-index"):
     ds = datasets.load_dataset(index_data, use_auth_token=auth_token)['train']
     index_data = {
@@ -77,20 +34,23 @@ def load_index(index_data: str = "clarin-knext/entity-linking-index"):
 
 
 def load_model(model_name: str = "clarin-knext/entity-linking-encoder"):
-    tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=auth_token)
-    model = AutoModel.from_pretrained(model_name, use_auth_token=auth_token)
-    return tokenizer, model
+    # tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=auth_token)
+    # model = AutoModel.from_pretrained(model_name, use_auth_token=auth_token)
+    model = pipeline("feature-extraction", model="clarin-knext/entity-linking-encoder", use_auth_token=auth_token)
+    # return tokenizer, model
+    return model
 
 
-tokenizer, model = load_model()
+# tokenizer, model = load_model()
+model = load_model()
 index = load_index()
 
 
 def predict(text: str = sample_text, top_k: int=3):
-    query = prepare_query(tokenizer, text)
+    # query = prepare_query(tokenizer, text)
     index_data, faiss_index = index
     # takes only the [CLS] embedding (for now)
-    query = model(query, return_tensors = "pt")[0][0].numpy().reshape(1, -1)  
+    query = model(query)[0][0].numpy().reshape(1, -1)  
     
     scores, indices = faiss_index.search(query, top_k)
     scores, indices = scores.tolist(), indices.tolist()