idk if i could do this less efficiently

app.py

@@ -153,7 +153,10 @@ def predict(queries, documents, max_characters) -> list[list[str]]:
     # Getting a structure like [[chunk, ...]]
     document_embeddings = [[] for _ in range(len(documents))]
     total_chars = 0
-    while total_chars < max_characters:
+    while (
+        total_chars < max_characters
+        and sum([len(x) for x in query_embeddings.values()]) > 0
+    ):
         for query, doc_scores in query_embeddings.items():
             if len(doc_scores) == 0:
                 continue
@@ -176,6 +179,12 @@ def predict(queries, documents, max_characters) -> list[list[str]]:
             document_embeddings[doc_idx].append(chunk_idx)
             total_chars += len(chunk)
 
+    # Get the actual text for the chunks
+    document_embeddings = [
+        [chunked_docs[doc_idx][chunk_idx] for chunk_idx in chunks]
+        for doc_idx, chunks in enumerate(document_embeddings)
+    ]
+
     return document_embeddings