kerfuffles

app.py

@@ -13,15 +13,6 @@ model = SentenceTransformer("Snowflake/snowflake-arctic-embed-m")
 model.to(device="cuda")
 
 
-def chunk(text, max_length=512):
-    chunks = []
-    while len(text) > max_length:
-        chunks.append(text[:max_length])
-        text = text[max_length:]
-    chunks.append(text)
-    return chunks
-
-
 @spaces.GPU
 def embed(queries, chunks) -> dict[str, list[tuple[str, float]]]:
     query_embeddings = model.encode(queries, prompt_name="query")
@@ -118,6 +109,15 @@ def convert(input_file) -> str:
     return convert_pandoc(input_file, input_file)
 
 
+def chunk_to_length(text, max_length=512):
+    chunks = []
+    while len(text) > max_length:
+        chunks.append(text[:max_length])
+        text = text[max_length:]
+    chunks.append(text)
+    return chunks
+
+
 @spaces.GPU
 def predict(queries, documents, max_characters) -> list[list[str]]:
     queries = queries.split("\n")
@@ -131,7 +131,7 @@ def predict(queries, documents, max_characters) -> list[list[str]]:
         return [[doc] for doc, _ in converted_docs]
 
     # Embed the documents in 512 character chunks
-    chunked_docs = [chunk(doc, 512) for doc in converted_docs]
+    chunked_docs = [chunk_to_length(doc, 512) for doc in converted_docs]
     embedded_docs = [embed(queries, chunks) for chunks in chunked_docs]
 
     # Get a structure like {query: [(doc_idx, chunk_idx, score), (doc_idx, chunk_idx, score), ...]}