an extended tokenizing function (as it was proposed in source project)

app.py

@@ -24,16 +24,16 @@ textbox = gr.Textbox(
 
 
 def prepare_query(tokenizer, query, max_seq_length=300):
-    # temporary solution
-    mention_start_token: str = "[unused0]"
-    mention_end_token: str = "[unused1]"
+    # temporary solution b/c of padding (which is unnecessary for inference)
+    start_token: str = "[unused0]"
+    end_token: str = "[unused1]"
     
-    left_context = query.split(mention_start_token)[0]
-    right_context = query.split(mention_end_token)[-1]
-    mention = query.split(mention_start_token)[-1].split(mention_end_token)[0]
+    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(
-        mention_start_token + mention + mention_end_token,
+        f"{start_token}{mention}{end_token}",
         add_special_tokens=False
     )['input_ids']
 
@@ -79,16 +79,15 @@ 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)
-    pipe = pipeline("feature-extraction", model=model, tokenizer=tokenizer)
-    return pipe
+    return tokenizer, model
 
 
-model = load_model()
+tokenizer, model = load_model()
 index = load_index()
 
 
 def predict(text: str = sample_text, top_k: int=3):
-    query = prepare_query(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)