Update chatpdf.py

chatpdf.py

@@ -383,14 +383,24 @@ class ChatPDF:
         """
         reference_results = []
         sim_contents = self.sim_model.most_similar(query, topn=self.similarity_top_k)
-        # Get reference results from corpus
-        hit_chunk_dict = dict()
-        for query_id, id_score_dict in sim_contents.items():
-            for corpus_id, s in id_score_dict.items():
+    
+        # Verificar si sim_contents es una lista
+        if isinstance(sim_contents, list):
+            # Ajusta esto dependiendo de la estructura de los elementos de la lista
+            for item in sim_contents:
+                # Ejemplo: si es una lista de tuplas (corpus_id, score)
+                corpus_id, _ = item
                 hit_chunk = self.sim_model.corpus[corpus_id]
                 reference_results.append(hit_chunk)
-                hit_chunk_dict[corpus_id] = hit_chunk
-
+    
+        elif isinstance(sim_contents, dict):
+            for query_id, id_score_dict in sim_contents.items():
+                for corpus_id, s in id_score_dict.items():
+                    hit_chunk = self.sim_model.corpus[corpus_id]
+                    reference_results.append(hit_chunk)
+        else:
+            logger.error(f"Unexpected type for sim_contents: {type(sim_contents)}")
+    
         if reference_results:
             if self.rerank_model is not None:
                 # Rerank reference results
@@ -410,7 +420,7 @@ class ChatPDF:
                         expanded_reference += self.sim_model.corpus.get(corpus_id + i + 1, '')
                     new_reference_results.append(expanded_reference)
                 reference_results = new_reference_results
-        return reference_results
+    return reference_results
 
     def predict_stream(
             self,