bug fixes

BBC-Esq · web-flow · commit 4f9e36e98dcd · 2025-05-04T10:54:10.000-04:00
diff --git a/src/choose_documents_and_vector_model.py b/src/choose_documents_and_vector_model.py
@@ -1,99 +1,147 @@
 from pathlib import Path
 import yaml
-from PySide6.QtWidgets import (QFileDialog, QDialog, QVBoxLayout, QTextEdit, 
-                              QPushButton, QHBoxLayout, QMessageBox)
-from create_symlinks import create_symlinks_parallel
-
-ALLOWED_EXTENSIONS = {'.pdf', '.docx', '.epub', '.txt', '.enex', '.eml', '.msg', '.csv', '.xls', '.xlsx', 
-                     '.rtf', '.odt', '.png', '.jpg', '.jpeg', '.bmp', '.gif', '.tif', '.tiff', '.html', 
-                     '.htm', '.md', '.doc'}
+from PySide6.QtCore import QThread, Signal, Qt, QElapsedTimer
+from PySide6.QtWidgets import (QFileDialog, QDialog, QVBoxLayout, QTextEdit, QPushButton, QHBoxLayout, QMessageBox, QProgressDialog, QApplication, QFileSystemModel)
+from multiprocessing import Pool, cpu_count
+from create_symlinks import _create_single_symlink
+ALLOWED_EXTENSIONS = {'.pdf', '.docx', '.epub', '.txt', '.enex', '.eml', '.msg', '.csv', '.xls', '.xlsx', '.rtf', '.odt', '.png', '.jpg', '.jpeg', '.bmp', '.gif', '.tif', '.tiff', '.html', '.htm', '.md', '.doc'}
 DOCS_FOLDER = "Docs_for_DB"
 CONFIG_FILE = "config.yaml"
-
+class SymlinkWorker(QThread):
+    progress = Signal(int)
+    finished = Signal(int, list)
+    def __init__(self, source, target_dir, parent=None):
+        super().__init__(parent)
+        self.source = source
+        self.target_dir = Path(target_dir)
+    def run(self):
+        if isinstance(self.source, (str, Path)):
+            dir_path = Path(self.source)
+            files = [str(p) for p in dir_path.iterdir() if p.is_file() and p.suffix.lower() in ALLOWED_EXTENSIONS]
+        else:
+            files = list(self.source)
+        total = len(files)
+        made = 0
+        errors = []
+        last_pct = -1
+        timer = QElapsedTimer()
+        timer.start()
+        step = max(1, total // 100) if total else 1
+        if total > 1000:
+            processes = min((total // 10000) + 1, cpu_count())
+            file_args = [(f, str(self.target_dir)) for f in files]
+            with Pool(processes=processes) as pool:
+                for i, (ok, err) in enumerate(pool.imap_unordered(_create_single_symlink, file_args), 1):
+                    if ok:
+                        made += 1
+                    if err:
+                        errors.append(err)
+                    if i % step == 0 or i == total:
+                        pct = int(i * 100 / total) if total else 100
+                        if pct != last_pct and timer.elapsed() > 500:
+                            self.progress.emit(pct)
+                            last_pct = pct
+                            timer.restart()
+        else:
+            for f in files:
+                if self.isInterruptionRequested():
+                    break
+                ok, err = _create_single_symlink((f, str(self.target_dir)))
+                if ok:
+                    made += 1
+                if err:
+                    errors.append(err)
+                if made % step == 0 or made == total:
+                    pct = int(made * 100 / total) if total else 100
+                    if pct != last_pct and timer.elapsed() > 500:
+                        self.progress.emit(pct)
+                        last_pct = pct
+                        timer.restart()
+        self.finished.emit(made, errors)
 def choose_documents_directory():
     current_dir = Path(__file__).parent.resolve()
     target_dir = current_dir / DOCS_FOLDER
     target_dir.mkdir(parents=True, exist_ok=True)
-
     msg_box = QMessageBox()
     msg_box.setWindowTitle("Selection Type")
     msg_box.setText("Would you like to select a directory or individual files?")
     dir_button = msg_box.addButton("Select Directory", QMessageBox.ActionRole)
     files_button = msg_box.addButton("Select Files", QMessageBox.ActionRole)
     cancel_button = msg_box.addButton("Cancel", QMessageBox.RejectRole)
-
     msg_box.exec()
     clicked_button = msg_box.clickedButton()
-
     if clicked_button == cancel_button:
         return
-
     file_dialog = QFileDialog()
-
+    def start_worker(source):
+        progress = QProgressDialog("Creating symlinks...", "Cancel", 0, 0)
+        progress.setWindowModality(Qt.WindowModal)
+        progress.setMinimumDuration(0)
+        worker = SymlinkWorker(source, target_dir)
+        main_window = _get_main_window()
+        if main_window and hasattr(main_window, 'databases_tab'):
+            db_tab = main_window.databases_tab
+            if hasattr(db_tab, 'docs_model') and db_tab.docs_model:
+                if hasattr(QFileSystemModel, 'DontWatchForChanges'):
+                    db_tab.docs_model.setOption(QFileSystemModel.DontWatchForChanges, True)
+                if hasattr(db_tab, 'docs_refresh'):
+                    db_tab.docs_refresh.start()
+        progress.canceled.connect(worker.requestInterruption)
+        def update_progress(pct):
+            if progress.maximum() == 0:
+                progress.setRange(0, 100)
+            progress.setValue(pct)
+        worker.progress.connect(update_progress)
+        def _done(count, errs):
+            if main_window and hasattr(main_window, 'databases_tab'):
+                db_tab = main_window.databases_tab
+                if hasattr(db_tab, 'docs_refresh'):
+                    db_tab.docs_refresh.stop()
+                if hasattr(db_tab, 'docs_model') and db_tab.docs_model:
+                    if hasattr(db_tab.docs_model, 'refresh'):
+                        db_tab.docs_model.refresh()
+                    elif hasattr(db_tab.docs_model, 'reindex'):
+                        db_tab.docs_model.reindex()
+                    if hasattr(QFileSystemModel, 'DontWatchForChanges'):
+                        db_tab.docs_model.setOption(QFileSystemModel.DontWatchForChanges, False)
+            progress.reset()
+            msg = f"Created {count} symlinks"
+            if errs:
+                msg += f" – {len(errs)} errors (see console)"
+                print(*errs, sep="\n")
+            QMessageBox.information(None, "Symlinks", msg)
+        worker.finished.connect(_done)
+        worker.progress.connect(update_progress)
+        worker.start()
+        choose_documents_directory._symlink_thread = worker
     if clicked_button == dir_button:
         file_dialog.setFileMode(QFileDialog.Directory)
         file_dialog.setOption(QFileDialog.ShowDirsOnly, True)
         selected_dir = file_dialog.getExistingDirectory(None, "Choose Directory for Database", str(current_dir))
         if selected_dir:
-            selected_dir_path = Path(selected_dir)
-            compatible_files = []
-            incompatible_files = []
-
-            for file_path in selected_dir_path.iterdir():
-                if file_path.is_file():
-                    if file_path.suffix.lower() in ALLOWED_EXTENSIONS:
-                        compatible_files.append(str(file_path))
-                    else:
-                        incompatible_files.append(file_path.name)
-
-            if incompatible_files:
-                if not show_incompatible_files_dialog(incompatible_files):
-                    return
-
-            if compatible_files:
-                try:
-                    count, errors = create_symlinks_parallel(compatible_files, target_dir)
-                    if errors:
-                        print("Errors occurred while creating symlinks:", errors)
-                except Exception as e:
-                    print(f"Error creating symlinks: {e}")
+            start_worker(Path(selected_dir))
     else:
         file_dialog.setFileMode(QFileDialog.ExistingFiles)
         file_paths = file_dialog.getOpenFileNames(None, "Choose Documents and Images for Database", str(current_dir))[0]
         if file_paths:
             compatible_files = []
             incompatible_files = []
-
             for file_path in file_paths:
                 path = Path(file_path)
                 if path.suffix.lower() in ALLOWED_EXTENSIONS:
                     compatible_files.append(str(path))
                 else:
                     incompatible_files.append(path.name)
-
-            if incompatible_files:
-                if not show_incompatible_files_dialog(incompatible_files):
-                    return
-
+            if incompatible_files and not show_incompatible_files_dialog(incompatible_files):
+                return
             if compatible_files:
-                try:
-                    count, errors = create_symlinks_parallel(compatible_files, target_dir)
-                    if errors:
-                        print("Errors occurred while creating symlinks:", errors)
-                except Exception as e:
-                    print(f"Error creating symlinks: {e}")
-
+                start_worker(compatible_files)
 def show_incompatible_files_dialog(incompatible_files):
-    dialog_text = (
-        "The following files cannot be added here due to their file extension:\n\n" +
-        "\n".join(incompatible_files) +
-        "\n\nHowever, if any of them are audio files you can still add them directly in the Tools Tab."
-        "\n\nClick 'Ok' to add the compatible documents only (remembering to add audio files separately) or 'Cancel' to back out completely."
-    )
+    dialog_text = ("The following files cannot be added here due to their file extension:\n\n" + "\n".join(incompatible_files) + "\n\nHowever, if any of them are audio files you can still add them directly in the Tools Tab."
+                   "\n\nClick 'Ok' to add the compatible documents only (remembering to add audio files separately) or 'Cancel' to back out completely.")
     incompatible_dialog = QDialog()
     incompatible_dialog.resize(800, 600)
     incompatible_dialog.setWindowTitle("Incompatible Files Detected")
-    
     layout = QVBoxLayout()
     text_edit = QTextEdit()
     text_edit.setReadOnly(True)
@@ -109,18 +157,19 @@ def show_incompatible_files_dialog(incompatible_files):
     ok_button.clicked.connect(incompatible_dialog.accept)
     cancel_button.clicked.connect(incompatible_dialog.reject)
     return incompatible_dialog.exec() == QDialog.Accepted
-
 def load_config():
     with open(CONFIG_FILE, 'r', encoding='utf-8') as stream:
         return yaml.safe_load(stream)
-
 def select_embedding_model_directory():
     initial_dir = Path('Models') if Path('Models').exists() else Path.home()
     chosen_directory = QFileDialog.getExistingDirectory(None, "Select Embedding Model Directory", str(initial_dir))
-
     if chosen_directory:
         config_file_path = Path(CONFIG_FILE)
         config_data = yaml.safe_load(config_file_path.read_text(encoding='utf-8')) if config_file_path.exists() else {}
         config_data["EMBEDDING_MODEL_NAME"] = chosen_directory
         config_file_path.write_text(yaml.dump(config_data), encoding='utf-8')
-        print(f"Selected directory: {chosen_directory}")
+def _get_main_window():
+    for widget in QApplication.topLevelWidgets():
+        if hasattr(widget, 'databases_tab'):
+            return widget
+    return None
diff --git a/src/create_symlinks.py b/src/create_symlinks.py
@@ -17,13 +17,13 @@ def create_symlinks_parallel(source: Union[str, Path, List[str], List[Path]],
                            target_dir: Union[str, Path] = "Docs_for_DB") -> Tuple[int, list]:
     """
     Create symbolic links using multiprocessing if the number of files exceeds 500.
-    
+
     Args:
         source: Can be either:
             - str or Path: Path to the source directory
             - List[str] or List[Path]: List of file paths
         target_dir: Path to the directory to store symlinks (default: 'Docs_for_DB')
-    
+
     Returns:
         tuple: (number of links created, list of errors)
     """
@@ -33,14 +33,12 @@ def create_symlinks_parallel(source: Union[str, Path, List[str], List[Path]],
         return 0, []
 
     try:
-        # Handle directory input
         if isinstance(source, (str, Path)) and not isinstance(source, list):
             source_dir = Path(source)
             if not source_dir.exists():
                 raise ValueError(f"Source directory does not exist: {source_dir}")
             files = [(str(p), str(target_dir)) for p in source_dir.iterdir() if p.is_file()]
 
-        # Handle list of files input
         elif isinstance(source, list):
             files = [(str(Path(p)), str(target_dir)) for p in source]
 
@@ -72,8 +70,8 @@ def create_symlinks_parallel(source: Union[str, Path, List[str], List[Path]],
             print("\nErrors occurred:")
             for error in errors:
                 print(error)
-                
+
         return count, errors
-        
+
     except Exception as e:
         raise RuntimeError(f"An error occurred: {str(e)}")
diff --git a/src/database_interactions.py b/src/database_interactions.py
@@ -53,7 +53,7 @@ def create(self):
         prepared_encode_kwargs = self.prepare_encode_kwargs()
         return HuggingFaceEmbeddings(
             model_name=self.model_name,
-            show_progress=not self.is_query,  # only show progress for database creation
+            show_progress=not self.is_query,
             model_kwargs=prepared_kwargs,
             encode_kwargs=prepared_encode_kwargs
         )
@@ -129,13 +129,6 @@ def prepare_kwargs(self):
         logging.debug(f"is_cuda: {is_cuda}")
         logging.debug(f"use_xformers: {use_xformers}")
 
-        # Add this tokenizer configuration to fix the error
-        stella_kwargs["tokenizer_kwargs"] = {
-            "max_length": 8192,
-            "padding": True,
-            "truncation": True
-        }
-
         stella_kwargs["config_kwargs"] = {
             "use_memory_efficient_attention": use_xformers,
             "unpad_inputs": use_xformers,
@@ -152,33 +145,31 @@ def prepare_kwargs(self):
 
 class AlibabaEmbedding(BaseEmbeddingModel):
     def prepare_kwargs(self):
-        logging.debug("Starting AlibabaEmbedding prepare_kwargs.")
         ali_kwargs = deepcopy(self.model_kwargs)
-        logging.debug(f"Original model_kwargs: {self.model_kwargs}")
-
-        compute_device = self.model_kwargs.get("device", "").lower()
+        compute_device = ali_kwargs.get("device", "").lower()
         is_cuda = compute_device == "cuda"
         use_xformers = is_cuda and supports_flash_attention()
-        logging.debug(f"Device: {compute_device}")
-        logging.debug(f"is_cuda: {is_cuda}")
-        logging.debug(f"use_xformers: {use_xformers}")
-
         ali_kwargs["tokenizer_kwargs"] = {
-            "max_length": 8192,
-            "padding": True,
-            "truncation": True
+            "padding": "longest",
+            "truncation": True,
+            "max_length": 8192
         }
-
         ali_kwargs["config_kwargs"] = {
             "use_memory_efficient_attention": use_xformers,
             "unpad_inputs": use_xformers,
             "attn_implementation": "eager" if use_xformers else "sdpa"
         }
-        logging.debug(f"Set 'config_kwargs': {ali_kwargs['config_kwargs']}")
-
-        logging.debug(f"Final ali_kwargs: {ali_kwargs}")
         return ali_kwargs
 
+    def prepare_encode_kwargs(self):
+        encode_kwargs = super().prepare_encode_kwargs()
+        encode_kwargs.update({
+            "padding": True,
+            "truncation": True,
+            "max_length": 8192
+        })
+        return encode_kwargs
+
 
 def create_vector_db_in_process(database_name):
     create_vector_db = CreateVectorDB(database_name=database_name)
@@ -306,7 +297,6 @@ def create_database(self, texts, embeddings):
             all_ids = []
             chunk_counters = defaultdict(int)
 
-            # Process all texts and generate IDs
             for doc in texts:
                 file_hash = doc.metadata.get('hash')
                 chunk_counters[file_hash] += 1
diff --git a/src/extract_metadata.py b/src/extract_metadata.py
@@ -1,4 +1,4 @@
-# works with langchain 0.3+
+# extract_metadata.py
 
 import os
 import datetime
diff --git a/src/gui_tabs_databases.py b/src/gui_tabs_databases.py

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-# works with langchain 0.3+`
	`1`	`+# extract_metadata.py`
`2`	`2`
`3`	`3`	`import os`
`4`	`4`	`import datetime`