use hf dataset

Dockerfile

@@ -6,15 +6,8 @@ RUN apt-get update && apt-get install -y \
     build-essential \
     curl \
     git \
-    unzip \
     && rm -rf /var/lib/apt/lists/*
 
-# Install AWS CLI
-RUN curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip" \
-    && unzip awscliv2.zip \
-    && ./aws/install \
-    && rm -rf aws awscliv2.zip
-
 # Install uv
 COPY --from=ghcr.io/astral-sh/uv:latest /uv /bin/uv
 

pyproject.toml

@@ -5,9 +5,9 @@ description = "Congressional Legislation Query and Analysis Tool"
 readme = "README.md"
 requires-python = ">=3.13"
 dependencies = [
-    "boto3>=1.35.0",
     "chromadb>=1.1.0",
     "datasets>=3.0.0",
+    "huggingface_hub>=0.19.0",
     "langchain>=0.3.27",
     "langchain-anthropic>=0.3.19",
     "langchain-chroma>=0.1.4",

src/legisqa_local/config/settings.py

@@ -38,7 +38,7 @@ def get_chroma_config():
     }
 
 def setup_chromadb():
-    """Setup ChromaDB - use persistent storage (/data) or download from S3 if needed"""
+    """Setup ChromaDB - use persistent storage (/data) or download from HF Dataset if needed"""
     logger.info("=== ChromaDB Setup Starting ===")
     
     chroma_config = get_chroma_config()
@@ -60,30 +60,27 @@ def setup_chromadb():
             logger.info(f"Updated CHROMA_PERSIST_DIRECTORY to: {persistent_chroma_path}")
             return persistent_chroma_path
         
-        # ChromaDB not found in persistent storage, try to download from S3
-        logger.info("ChromaDB not found in persistent storage, checking S3 configuration...")
+        # ChromaDB not found in persistent storage, try to download from HF Dataset
+        logger.info("ChromaDB not found in persistent storage, checking HF Dataset configuration...")
         
-        s3_bucket = os.getenv("CHROMA_S3_BUCKET", "")
-        s3_prefix = os.getenv("CHROMA_S3_PREFIX", "")
+        dataset_repo = os.getenv("CHROMA_DATASET_REPO", "hyperdemocracy/usc-chroma-vecs-v1-chunks-v1-s8192-o512-sentence-transformers-static-retrieval-mrl-en-v1")
+        logger.info(f"HF Dataset repo: {dataset_repo}")
         
-        logger.info(f"S3 Bucket: {s3_bucket}")
-        logger.info(f"S3 Prefix: {s3_prefix}")
-        
-        if s3_bucket and s3_prefix:
-            logger.info(f"📥 Downloading ChromaDB from S3 to persistent storage...")
-            logger.info(f"   Source: s3://{s3_bucket}/{s3_prefix}")
+        if dataset_repo:
+            logger.info(f"📥 Downloading ChromaDB from HF Dataset to persistent storage...")
+            logger.info(f"   Source: {dataset_repo}")
             logger.info(f"   Target: {persistent_chroma_path}")
             
-            success = download_chromadb_from_s3(s3_bucket, s3_prefix, persistent_chroma_path)
+            success = download_chromadb_from_hf_dataset(dataset_repo, persistent_chroma_path)
             if success:
                 # Update environment variable to point to persistent storage
                 os.environ["CHROMA_PERSIST_DIRECTORY"] = persistent_chroma_path
                 logger.info(f"✅ ChromaDB download successful! Updated path to: {persistent_chroma_path}")
                 return persistent_chroma_path
             else:
-                logger.error("❌ ChromaDB download from S3 failed!")
+                logger.error("❌ ChromaDB download from HF Dataset failed!")
         else:
-            logger.error("❌ No S3 configuration found (CHROMA_S3_BUCKET, CHROMA_S3_PREFIX)")
+            logger.error("❌ No HF Dataset configuration found (CHROMA_DATASET_REPO)")
             logger.info("Available environment variables:")
             for key, value in os.environ.items():
                 if "CHROMA" in key:
@@ -106,92 +103,58 @@ def setup_chromadb():
     logger.info("=== ChromaDB Setup Complete ===")
     return chroma_path
 
-def download_chromadb_from_s3(bucket: str, prefix: str, local_path: str) -> bool:
-    """Download ChromaDB from S3"""
-    logger.info(f"Starting S3 download: s3://{bucket}/{prefix} -> {local_path}")
+def download_chromadb_from_hf_dataset(dataset_repo: str, local_path: str) -> bool:
+    """Download ChromaDB from HuggingFace Dataset"""
+    logger.info(f"Starting HF Dataset download: {dataset_repo} -> {local_path}")
     
     try:
-        import subprocess
+        from huggingface_hub import snapshot_download
         import os
         
         # Ensure target directory exists
         logger.info(f"Creating target directory: {local_path}")
         os.makedirs(local_path, exist_ok=True)
         
-        # Use AWS CLI to sync from S3 (no credentials needed for public buckets)
-        s3_url = f"s3://{bucket}/{prefix}"
-        cmd = ["aws", "s3", "sync", s3_url, local_path, "--no-sign-request"]
-        
-        logger.info(f"Running AWS CLI command: {' '.join(cmd)}")
-        result = subprocess.run(cmd, capture_output=True, text=True)
-        
-        if result.returncode == 0:
-            logger.info("✅ ChromaDB download from S3 (AWS CLI) complete!")
-            if result.stdout:
-                logger.info(f"AWS CLI output: {result.stdout}")
+        # Download the dataset using snapshot_download
+        logger.info(f"Downloading dataset: {dataset_repo}")
+        logger.info("This may take several minutes for large datasets...")
+        
+        downloaded_path = snapshot_download(
+            repo_id=dataset_repo,
+            repo_type="dataset",
+            local_dir=local_path,
+            cache_dir="/tmp/hf_chromadb_cache",
+            resume_download=True,  # Resume if interrupted
+            local_dir_use_symlinks=False  # Copy files instead of symlinks
+        )
+        
+        logger.info(f"✅ ChromaDB download from HF Dataset complete!")
+        logger.info(f"Downloaded to: {downloaded_path}")
+        
+        # Verify the download by checking for expected files
+        if os.path.exists(local_path) and os.listdir(local_path):
+            file_count = sum(len(files) for _, _, files in os.walk(local_path))
+            total_size = sum(
+                os.path.getsize(os.path.join(dirpath, filename))
+                for dirpath, _, filenames in os.walk(local_path)
+                for filename in filenames
+            ) / (1024 * 1024 * 1024)  # Convert to GB
+            
+            logger.info(f"📊 Download verification:")
+            logger.info(f"   Files: {file_count}")
+            logger.info(f"   Total size: {total_size:.2f} GB")
+            
             return True
         else:
-            logger.error(f"❌ AWS CLI failed with return code {result.returncode}")
-            logger.error(f"AWS CLI stderr: {result.stderr}")
-            if result.stdout:
-                logger.error(f"AWS CLI stdout: {result.stdout}")
+            logger.error("❌ Download completed but no files found in target directory")
             return False
             
-    except FileNotFoundError:
-        logger.warning("❌ AWS CLI not found. Trying with boto3...")
-        return download_chromadb_from_s3_boto3(bucket, prefix, local_path)
-    except Exception as e:
-        logger.error(f"❌ Error downloading from S3: {e}")
-        return False
-
-def download_chromadb_from_s3_boto3(bucket: str, prefix: str, local_path: str) -> bool:
-    """Download ChromaDB from S3 using boto3 (fallback)"""
-    try:
-        import boto3
-        from botocore import UNSIGNED
-        from botocore.config import Config
-        import os
-        
-        logger.info("📦 Using boto3 for S3 download...")
-        
-        # Create S3 client with no credentials (for public buckets)
-        s3 = boto3.client('s3', config=Config(signature_version=UNSIGNED))
-        
-        # List objects in the S3 prefix
-        logger.info(f"Listing objects in s3://{bucket}/{prefix}")
-        paginator = s3.get_paginator('list_objects_v2')
-        pages = paginator.paginate(Bucket=bucket, Prefix=prefix)
-        
-        os.makedirs(local_path, exist_ok=True)
-        
-        file_count = 0
-        for page in pages:
-            if 'Contents' in page:
-                for obj in page['Contents']:
-                    key = obj['Key']
-                    # Get relative path by removing prefix
-                    relative_path = key[len(prefix):].lstrip('/')
-                    if relative_path:  # Skip empty paths
-                        local_file_path = os.path.join(local_path, relative_path)
-                        
-                        # Create directory if needed
-                        os.makedirs(os.path.dirname(local_file_path), exist_ok=True)
-                        
-                        # Download file
-                        file_count += 1
-                        if file_count % 10 == 0:
-                            logger.info(f"Downloaded {file_count} files...")
-                        
-                        s3.download_file(bucket, key, local_file_path)
-        
-        logger.info(f"✅ ChromaDB download from S3 (boto3) complete! Downloaded {file_count} files.")
-        return True
-        
     except ImportError:
-        logger.error("❌ boto3 not available. Please install: pip install boto3")
+        logger.error("❌ huggingface_hub not available. Please install: pip install huggingface_hub")
         return False
     except Exception as e:
-        logger.error(f"❌ Error downloading from S3 with boto3: {e}")
+        logger.error(f"❌ Error downloading from HF Dataset: {e}")
+        logger.error(f"Exception type: {type(e).__name__}")
         return False
 
 # Embedding model configuration