Fix

app.py

@@ -1,13 +1,11 @@
 import os
 import json
 
-import spaces
 import gradio
 
 import numpy
 import pandas
 
-from transformers import AutoModelForCausalLM, AutoTokenizer
 
 import pyparseit
 
@@ -35,50 +33,14 @@ model_options = [
 
 
 from src.SemanticSearch import SemanticSearch
+from src.Inference import model_completion
 
 extractor = SemanticSearch()
 extractor.load_ne_from_kg(SPARQL_ENDPOINT)
 extractor.build_vector_db()
 extractor.load_vector_db()
 
-@spaces.GPU
-def model_completion(messages, model_name, model_temperature, model_thinking):
 
-    # load the tokenizer and the model
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        torch_dtype="auto",
-        device_map="auto"
-    )
-
-    text = tokenizer.apply_chat_template(
-        messages,
-        tokenize=False,
-        add_generation_prompt=True,
-        enable_thinking=model_thinking
-    )
-
-    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
-
-    sample = True
-
-    if model_temperature == 0:
-        sample = False
-
-
-    # conduct text completion
-    generated_ids = model.generate(
-        **model_inputs,
-        max_new_tokens=4096,
-        do_sample=sample,
-        temperature=model_temperature
-    )
-    output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist() 
-
-    content = tokenizer.decode(output_ids, skip_special_tokens=True).strip("\n")
-
-    return content
 
 
 def sparql_json_to_df(sparql_json):

src/Inference.py

@@ -0,0 +1,43 @@
+import spaces
+
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+@spaces.GPU
+def model_completion(messages, model_name, model_temperature, model_thinking):
+
+    # load the tokenizer and the model
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype="auto",
+        device_map="auto"
+    )
+
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
+        enable_thinking=model_thinking
+    )
+
+    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+
+    sample = True
+
+    if model_temperature == 0:
+        sample = False
+
+
+    # conduct text completion
+    generated_ids = model.generate(
+        **model_inputs,
+        max_new_tokens=4096,
+        do_sample=sample,
+        temperature=model_temperature
+    )
+    output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist() 
+
+    content = tokenizer.decode(output_ids, skip_special_tokens=True).strip("\n")
+
+    return content

src/SemanticSearch.py

@@ -30,9 +30,17 @@ WHERE {
   #FILTER(lang(?ne_label) = "en" || lang(?ne_label) = "")
   #FILTER(lang(?class_label) = "en" || lang(?class_label) = "")
 }
-LIMIT 128
 """
 
+# HF seems to use 3.10!
+def batched(iterable, n):
+    if n < 1:
+        raise ValueError('n must be at least one')
+    it = iter(iterable)
+
+    while batch := tuple(itertools.islice(it, n)):
+        yield batch
+
 class SemanticSearch:
 
     def __init__(self, embeddings_model="BAAI/bge-base-en-v1.5", reranking_model="BAAI/bge-reranker-v2-m3"):
@@ -174,7 +182,7 @@ class SemanticSearch:
 
         print("Got ", len(documents), "sentences")
 
-        for sentences_batch in tqdm.tqdm(list(itertools.batched(documents, 512)), desc="Generating embeddings"):
+        for sentences_batch in tqdm.tqdm(list(batched(documents, 512)), desc="Generating embeddings"):
 
             embeddings += self.get_text_embeddings_local(sentences_batch)