Fix module import issues and data loading bugs

download_from_hub.py

@@ -4,7 +4,6 @@ import pickle
 import sys
 import numpy as np
 from huggingface_hub import hf_hub_download, list_repo_files
-from datasets import load_dataset
 
 def ensure_dirs():
     """Create necessary directories if they don't exist."""
@@ -63,6 +62,7 @@ def download_datasets():
     
     # Download embeddings
     try:
+        from datasets import load_dataset
         print("Downloading embeddings...")
         # First check what files are available in the dataset repository
         try:
@@ -124,6 +124,7 @@ def download_datasets():
     
     # Download document chunks
     try:
+        from datasets import load_dataset
         print("Downloading document chunks...")
         # First check what files are available
         try:

embedder.py

@@ -0,0 +1,90 @@
+import time
+import numpy as np
+from tqdm import tqdm
+from openai import OpenAI
+from typing import List, Dict, Any, Optional
+import os
+
+# Get API key from environment variable
+OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "")
+EMBEDDING_MODEL = "text-embedding-3-small"
+EMBEDDING_BATCH_SIZE = 10
+
+class TextEmbedder:
+    """Class for generating embeddings for document chunks using OpenAI's embeddings API."""
+    
+    def __init__(self, model: str = EMBEDDING_MODEL, batch_size: int = EMBEDDING_BATCH_SIZE):
+        """
+        Initialize the TextEmbedder with the specified embedding model and batch size.
+        
+        Args:
+            model: The OpenAI embedding model to use
+            batch_size: Number of chunks to embed per API call
+        """
+        self.model = model
+        self.batch_size = batch_size
+        self.client = OpenAI(api_key=OPENAI_API_KEY)
+        self.embedding_dim = 1536  # Default dimension for text-embedding-3-small
+    
+    def get_embedding_for_text(self, text: str) -> List[float]:
+        """Generate embedding for a single text."""
+        try:
+            response = self.client.embeddings.create(
+                input=[text],
+                model=self.model
+            )
+            return response.data[0].embedding
+        except Exception as e:
+            print(f"Error generating embedding: {e}")
+            return [0.0] * self.embedding_dim
+    
+    def get_embeddings_for_texts(self, texts: List[str]) -> List[List[float]]:
+        """
+        Compute embeddings for a list of texts using batched API calls.
+        
+        Args:
+            texts: List of text chunks to embed
+            
+        Returns:
+            List of embedding vectors
+        """
+        embeddings = []
+        for i in tqdm(range(0, len(texts), self.batch_size), desc="Embedding chunks"):
+            batch = texts[i:i + self.batch_size]
+            try:
+                response = self.client.embeddings.create(
+                    input=batch,
+                    model=self.model
+                )
+                # Extract embeddings from the response
+                for item in response.data:
+                    embeddings.append(item.embedding)
+            except Exception as e:
+                print(f"Error embedding batch starting at index {i}: {e}")
+                # Append placeholder zero vectors for failed texts
+                for _ in batch:
+                    embeddings.append([0.0] * self.embedding_dim)
+            # Brief pause to avoid rate limits
+            time.sleep(0.2)
+        
+        return embeddings
+    
+    def get_query_embedding(self, query: str) -> np.ndarray:
+        """
+        Generate embedding for a query string and return as numpy array.
+        
+        Args:
+            query: The query text to embed
+            
+        Returns:
+            Numpy array of the embedding
+        """
+        try:
+            q_response = self.client.embeddings.create(
+                input=[query], 
+                model=self.model
+            )
+            return np.array(q_response.data[0].embedding, dtype='float32').reshape(1, -1)
+        except Exception as e:
+            print(f"Error creating embedding for query: {e}")
+            return np.zeros((1, self.embedding_dim), dtype='float32') 

src/__init__.py

@@ -1,13 +1,4 @@
 """Agentic Defensor - An agentic RAG system for legal defense analysis."""
 
-# Ensure subpackages are properly importable
-try:
-    from . import embeddings
-    from . import models
-    from . import agents
-    from . import utils
-    from . import data
-    from . import api
-except ImportError as e:
-    import sys
-    print(f"Warning: Not all subpackages could be imported: {e}", file=sys.stderr) 
+# No explicit imports to avoid circular dependencies
+# Each module will be imported directly where needed 

src/models/retriever.py

@@ -4,7 +4,19 @@ import numpy as np
 from typing import List, Dict, Any, Tuple, Optional
 
 from src.utils.config import TOP_K, FAISS_INDEX_PATH, DOC_CHUNKS_PATH
-from src.embeddings.embedder import TextEmbedder
+
+# Try to import from the proper location, otherwise use the local copy
+try:
+    from src.embeddings.embedder import TextEmbedder
+except ImportError:
+    try:
+        import sys
+        import os
+        sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
+        from embedder import TextEmbedder
+        print("Using local copy of embedder.py")
+    except ImportError as e:
+        print(f"Error importing TextEmbedder: {e}")
 
 class Retriever:
     """