Update chatpdf.py

chatpdf.py

@@ -420,13 +420,28 @@ 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():
-                hit_chunk = self.sim_model.corpus[corpus_id]
-                reference_results.append(hit_chunk)
-                hit_chunk_dict[corpus_id] = hit_chunk
+    
+        # Verificar si sim_contents es una lista o un diccionario
+        if isinstance(sim_contents, list):
+            for item in sim_contents:
+                # Ajustar según la estructura real de item
+                corpus_id = item[0] if isinstance(item, (list, tuple)) else item  # Asegurarse de que corpus_id sea el valor correcto
+                if isinstance(corpus_id, dict):
+                    # Extraer el valor necesario si corpus_id es un diccionario
+                    corpus_id = next(iter(corpus_id.keys()))  # Tomar la primera clave como ejemplo
+                if corpus_id in self.sim_model.corpus:
+                    hit_chunk = self.sim_model.corpus[corpus_id]
+                    reference_results.append(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():
+                    if corpus_id in self.sim_model.corpus:
+                        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: