refactor: adjust splade script

FlorianWoelki · Oct 31, 2024 · cffaa5a · cffaa5a
1 parent caba1af
commit cffaa5a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/scripts/splade.py b/scripts/splade.py
@@ -36,7 +36,7 @@ def clean_text(text):
 for example in tqdm(training_dataset, desc="Processing"):
     text = clean_text(example['Overview'])
     texts.append(text)
-    if len(texts) == 11: # temporary
+    if len(texts) == 25000: # temporary
         break
 
 def sparse_vector_to_dict(vector):
@@ -65,7 +65,7 @@ def compute_similarity(vec1, vec2):
 for text in tqdm(texts, desc="Generating sparse vectors"):
     sparse_vectors.append(sparse_vector_to_dict(create_splade(text)))
 
-num_queries = 2
+num_queries = round(len(texts) / 10) # Like in `generator_sparse.rs` file
 queries = []
 query_vectors = []
 groundtruth = []