minor fixes

document_classifier.py

@@ -1,7 +1,6 @@
 from langchain_community.document_loaders import PyPDFLoader
 from transformers import pipeline
 import torch
-from concurrent.futures import ThreadPoolExecutor, as_completed
 from collections import defaultdict
 import time
 
@@ -23,12 +22,10 @@ class DocumentClassifier:
         self, 
         pages_per_group=2,
         min_confidence=0.35,
-        max_workers=4,
         model_name="cross-encoder/nli-deberta-v3-small"
     ):
         self.pages_per_group = pages_per_group
         self.min_confidence = min_confidence
-        self.max_workers = max_workers
         self.model_name = model_name
         self.classifier = None
         
@@ -111,25 +108,31 @@ class DocumentClassifier:
     
     def _classify_groups_parallel(self, groups):
         results = []
-        
-        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
-            future_to_group = {
-                executor.submit(self._classify_single_group, group): group
-                for group in groups
-            }
-            
-            for future in as_completed(future_to_group):
-                group = future_to_group[future]
-                try:
-                    result = future.result()
-                    result['page_numbers'] = group['page_numbers']
-                    results.append(result)
-                except Exception as e:
-                    print(f"[Classifier] Group classification failed: {e}")
-        
+        texts = [g['text'] for g in groups]
+
+        # Use pipeline's native batching — faster than ThreadPoolExecutor,
+        # especially on GPU, and avoids thread-safety issues with PyTorch.
+        batch_results = self.classifier(texts, self.LABELS, multi_label=True, batch_size=8)
+
+        for group, result in zip(groups, batch_results):
+            primary_type = result['labels'][0]
+            primary_score = result['scores'][0]
+
+            if primary_score < self.min_confidence:
+                primary_type = 'other'
+
+            scores = {label: score for label, score in zip(result['labels'], result['scores'])}
+            results.append({
+                'type': primary_type,
+                'confidence': primary_score,
+                'scores': scores,
+                'page_numbers': group['page_numbers']
+            })
+
         return results
     
     def _classify_single_group(self, group):
+        # Kept for single-group use if needed directly
         text = group['text']
         
         if not text.strip():
@@ -139,6 +142,9 @@ class DocumentClassifier:
         
         primary_type = result['labels'][0]
         primary_score = result['scores'][0]
+
+        if primary_score < self.min_confidence:
+            primary_type = 'other'
         
         scores = {
             label: score 

rag_setup.py

@@ -18,7 +18,6 @@ class RAG_Setup:
         self.classifier = DocumentClassifier(
             pages_per_group=2,
             min_confidence=0.35,
-            max_workers=4,
             model_name="cross-encoder/nli-deberta-v3-small"
         )
 

requirements.txt

@@ -8,5 +8,4 @@ transformers
 sentence-transformers
 torch
 pypdf
-gradio
-pyaudioop
+gradio

user_data.py

@@ -25,12 +25,6 @@ def initialize_db():
             created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
             FOREIGN KEY (user_id) REFERENCES users(id)
         );
-
-        CREATE TABLE IF NOT EXISTS document_classifications (
-            file_hash TEXT PRIMARY KEY,
-            doc_type TEXT NOT NULL,
-            classified_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
-        );
     """)
     conn.commit()
     conn.close()
@@ -57,32 +51,3 @@ def user_exists(user_id):
     exists = cursor.fetchone() is not None
     conn.close()
     return exists
-
-
-def get_document_label(file_hash: str):
-    conn = get_connection()
-    cursor = conn.cursor()
-    cursor.execute(
-        "SELECT doc_type FROM document_classifications WHERE file_hash = ?",
-        (file_hash,)
-    )
-    row = cursor.fetchone()
-    conn.close()
-    return row["doc_type"] if row else None
-
-
-def save_document_label(file_hash: str, doc_type: str):
-    conn = get_connection()
-    cursor = conn.cursor()
-    cursor.execute(
-        """
-        INSERT INTO document_classifications (file_hash, doc_type)
-        VALUES (?, ?)
-        ON CONFLICT(file_hash) DO UPDATE SET
-            doc_type = excluded.doc_type,
-            classified_at = CURRENT_TIMESTAMP
-        """,
-        (file_hash, doc_type)
-    )
-    conn.commit()
-    conn.close()