Update chatpdf.py

chatpdf.py

@@ -373,44 +373,44 @@ class ChatPDF:
         return scores
 
     def get_reference_results(self, query: str):
-    reference_results = []
-    sim_contents = self.sim_model.most_similar(query, topn=self.similarity_top_k)
-    
-    # Ajustar según el tipo de retorno de sim_contents
-    if isinstance(sim_contents, dict):  # Si es un diccionario
+        """
+        Get reference results.
+            1. Similarity model get similar chunks
+            2. Rerank similar chunks
+            3. Expand reference context chunk
+        :param query:
+        :return:
+        """
+        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)
-    elif isinstance(sim_contents, list):  # Si es una lista
-        for item in sim_contents:
-            # Ajusta esto dependiendo de la estructura de los elementos de la lista
-            # Ejemplo: si es una lista de (corpus_id, score) tuplas
-            corpus_id, _ = item
-            hit_chunk = self.sim_model.corpus[corpus_id]
-            reference_results.append(hit_chunk)
-
-    # Resto del código...
-    if reference_results:
-        if self.rerank_model is not None:
-            # Rerank reference results
-            rerank_scores = self._get_reranker_score(query, reference_results)
-            logger.debug(f"rerank_scores: {rerank_scores}")
-            # Get rerank top k chunks
-            reference_results = [reference for reference, score in sorted(
-                zip(reference_results, rerank_scores), key=lambda x: x[1], reverse=True)][:self.rerank_top_k]
-            hit_chunk_dict = {corpus_id: hit_chunk for corpus_id, hit_chunk in hit_chunk_dict.items() if
-                              hit_chunk in reference_results}
-        # Expand reference context chunk
-        if self.num_expand_context_chunk > 0:
-            new_reference_results = []
-            for corpus_id, hit_chunk in hit_chunk_dict.items():
-                expanded_reference = self.sim_model.corpus.get(corpus_id - 1, '') + hit_chunk
-                for i in range(self.num_expand_context_chunk):
-                    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
+                hit_chunk_dict[corpus_id] = hit_chunk
+
+        if reference_results:
+            if self.rerank_model is not None:
+                # Rerank reference results
+                rerank_scores = self._get_reranker_score(query, reference_results)
+                logger.debug(f"rerank_scores: {rerank_scores}")
+                # Get rerank top k chunks
+                reference_results = [reference for reference, score in sorted(
+                    zip(reference_results, rerank_scores), key=lambda x: x[1], reverse=True)][:self.rerank_top_k]
+                hit_chunk_dict = {corpus_id: hit_chunk for corpus_id, hit_chunk in hit_chunk_dict.items() if
+                                  hit_chunk in reference_results}
+            # Expand reference context chunk
+            if self.num_expand_context_chunk > 0:
+                new_reference_results = []
+                for corpus_id, hit_chunk in hit_chunk_dict.items():
+                    expanded_reference = self.sim_model.corpus.get(corpus_id - 1, '') + hit_chunk
+                    for i in range(self.num_expand_context_chunk):
+                        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
 
     def predict_stream(
             self,