Update chatpdf.py

chatpdf.py

@@ -139,7 +139,7 @@ class ChatPDF:
             self,
             similarity_model: SimilarityABC = None,
             generate_model_type: str = "auto",
-            generate_model_name_or_path: str = "Qwen/Qwen2-0.5B-Instruct",
+            generate_model_name_or_path: str = "LenguajeNaturalAI/leniachat-qwen2-1.5B-v0",
             lora_model_name_or_path: str = None,
             corpus_files: Union[str, List[str]] = None,
             save_corpus_emb_dir: str = "corpus_embs/",
@@ -188,7 +188,7 @@ class ChatPDF:
         if similarity_model is not None:
             self.sim_model = similarity_model
         else:
-            m1 = BertSimilarity(model_name_or_path="shibing624/text2vec-base-multilingual", device=self.device)
+            m1 = BertSimilarity(model_name_or_path="jaimevera1107/all-MiniLM-L6-v2-similarity-es", device=self.device)
             m2 = BM25Similarity()
             default_sim_model = EnsembleSimilarity(similarities=[m1, m2], weights=[0.5, 0.5], c=2)
             self.sim_model = default_sim_model
@@ -420,28 +420,13 @@ class ChatPDF:
         """
         reference_results = []
         sim_contents = self.sim_model.most_similar(query, topn=self.similarity_top_k)
-    
-        # 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)}")
-
+        # 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
 
         if reference_results:
             if self.rerank_model is not None:
@@ -579,9 +564,9 @@ class ChatPDF:
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument("--sim_model_name", type=str, default="shibing624/text2vec-base-multilingual")
+    parser.add_argument("--sim_model_name", type=str, default="jaimevera1107/all-MiniLM-L6-v2-similarity-es")
     parser.add_argument("--gen_model_type", type=str, default="auto")
-    parser.add_argument("--gen_model_name", type=str, default="Qwen/Qwen2-0.5B-Instruct")
+    parser.add_argument("--gen_model_name", type=str, default="LenguajeNaturalAI/leniachat-qwen2-1.5B-v0")
     parser.add_argument("--lora_model", type=str, default=None)
     parser.add_argument("--rerank_model_name", type=str, default="")
     parser.add_argument("--corpus_files", type=str, default="Acuerdo009.pdf")