Update src/kaggle_loader.py

src/kaggle_loader.py

@@ -0,0 +1,240 @@
+import os
+import pandas as pd
+import json
+from typing import List, Optional
+from langchain_core.documents import Document
+from langchain_community.document_loaders import CSVLoader, JSONLoader
+import kaggle
+
+class KaggleDataLoader:
+    """Load and process Kaggle datasets for RAG."""
+    
+    def __init__(self, kaggle_username: Optional[str] = None, kaggle_key: Optional[str] = None):
+        """
+        Initialize Kaggle loader.
+        
+        Args:
+            kaggle_username: Your Kaggle username (optional if using kaggle.json)
+            kaggle_key: Your Kaggle API key (optional if using kaggle.json)
+        """
+        self.kaggle_username = kaggle_username
+        self.kaggle_key = kaggle_key
+        
+        # Try to load credentials from kaggle.json first
+        self._load_kaggle_credentials()
+        
+        # Set Kaggle credentials (either from kaggle.json or parameters)
+        if self.kaggle_username and self.kaggle_key:
+            os.environ['KAGGLE_USERNAME'] = self.kaggle_username
+            os.environ['KAGGLE_KEY'] = self.kaggle_key
+            print("Kaggle credentials loaded successfully")
+        else:
+            print("Warning: No Kaggle credentials found. Please set up kaggle.json or provide credentials.")
+    
+    def _load_kaggle_credentials(self):
+        """Load Kaggle credentials from kaggle.json file."""
+        # Common locations for kaggle.json
+        possible_paths = [
+            os.path.expanduser("~/.kaggle/kaggle.json"),
+            os.path.expanduser("~/kaggle.json"),
+            "./kaggle.json",
+            os.path.join(os.getcwd(), "kaggle.json")
+        ]
+        
+        for path in possible_paths:
+            if os.path.exists(path):
+                try:
+                    with open(path, 'r') as f:
+                        credentials = json.load(f)
+                    
+                    # Extract username and key from kaggle.json
+                    if 'username' in credentials and 'key' in credentials:
+                        self.kaggle_username = credentials['username']
+                        self.kaggle_key = credentials['key']
+                        print(f"Loaded Kaggle credentials from {path}")
+                        return
+                    else:
+                        print(f"Invalid kaggle.json format at {path}. Expected 'username' and 'key' fields.")
+                        
+                except Exception as e:
+                    print(f"Error reading kaggle.json from {path}: {e}")
+        
+        print("No valid kaggle.json found in common locations:")
+        for path in possible_paths:
+            print(f"  - {path}")
+        print("Please create kaggle.json with your Kaggle API credentials.")
+    
+    def download_dataset(self, dataset_name: str, download_path: str = "./data") -> str:
+        """
+        Download a Kaggle dataset.
+        
+        Args:
+            dataset_name: Dataset name in format 'username/dataset-name'
+            download_path: Where to save the dataset
+            
+        Returns:
+            Path to downloaded dataset
+        """
+        if not self.kaggle_username or not self.kaggle_key:
+            raise ValueError("Kaggle credentials not found. Please set up kaggle.json or provide credentials.")
+        
+        try:
+            # Create a unique directory for this dataset
+            dataset_dir = dataset_name.replace('/', '_')
+            full_download_path = os.path.join(download_path, dataset_dir)
+            
+            # Create the directory if it doesn't exist
+            os.makedirs(full_download_path, exist_ok=True)
+            
+            kaggle.api.authenticate()
+            kaggle.api.dataset_download_files(dataset_name, path=full_download_path, unzip=True)
+            print(f"Dataset {dataset_name} downloaded successfully to {full_download_path}")
+            return full_download_path
+        except Exception as e:
+            print(f"Error downloading dataset: {e}")
+            raise
+    
+    def load_csv_dataset(self, file_path: str, text_columns: List[str], chunk_size: int = 100) -> List[Document]:
+        """Load documents from a CSV file."""
+        try:
+            df = pd.read_csv(file_path)
+            documents = []
+            
+            # For FAQ datasets, try to combine question and answer columns
+            if 'Questions' in df.columns and 'Answers' in df.columns:
+                print(f"Processing FAQ dataset with {len(df)} Q&A pairs")
+                for idx, row in df.iterrows():
+                    question = str(row['Questions']).strip()
+                    answer = str(row['Answers']).strip()
+                    
+                    # Create a document with question prominently featured for better retrieval
+                    content = f"QUESTION: {question}\n\nANSWER: {answer}"
+                    documents.append(Document(
+                        page_content=content,
+                        metadata={"source": file_path, "type": "faq", "question_id": idx, "question": question}
+                    ))
+            else:
+                # Fallback to original method for other CSV files
+                print(f"Processing regular CSV with columns: {text_columns}")
+                for idx, row in df.iterrows():
+                    # Combine specified text columns
+                    text_parts = []
+                    for col in text_columns:
+                        if col in df.columns and pd.notna(row[col]):
+                            text_parts.append(str(row[col]).strip())
+                    
+                    if text_parts:
+                        content = " ".join(text_parts)
+                        documents.append(Document(
+                            page_content=content,
+                            metadata={"source": file_path, "row": idx}
+                        ))
+            
+            print(f"Created {len(documents)} documents from CSV")
+            return documents
+            
+        except Exception as e:
+            print(f"Error loading CSV dataset: {e}")
+            return []
+    
+    def load_json_dataset(self, file_path: str, text_field: str = "text",
+                         metadata_fields: Optional[List[str]] = None) -> List[Document]:
+        """
+        Load JSON data and convert to documents.
+        
+        Args:
+            file_path: Path to JSON file
+            text_field: Field name containing the main text
+            metadata_fields: Fields to include as metadata
+            
+        Returns:
+            List of Document objects
+        """
+        with open(file_path, 'r') as f:
+            data = json.load(f)
+        
+        documents = []
+        
+        for item in data:
+            text_content = item.get(text_field, "")
+            
+            # Create metadata
+            metadata = {"source": file_path}
+            if metadata_fields:
+                for field in metadata_fields:
+                    if field in item:
+                        metadata[field] = item[field]
+            
+            documents.append(Document(
+                page_content=text_content,
+                metadata=metadata
+            ))
+        
+        return documents
+    
+    def load_text_dataset(self, file_path: str, chunk_size: int = 1000) -> List[Document]:
+        """
+        Load plain text data and convert to documents.
+        
+        Args:
+            file_path: Path to text file
+            chunk_size: Number of characters per document
+            
+        Returns:
+            List of Document objects
+        """
+        with open(file_path, 'r', encoding='utf-8') as f:
+            text = f.read()
+        
+        documents = []
+        
+        for i in range(0, len(text), chunk_size):
+            chunk = text[i:i+chunk_size]
+            
+            documents.append(Document(
+                page_content=chunk,
+                metadata={
+                    "source": file_path,
+                    "chunk_id": i // chunk_size,
+                    "start_char": i,
+                    "end_char": min(i + chunk_size, len(text))
+                }
+            ))
+        
+        return documents
+
+# Example usage functions
+def load_kaggle_csv_example():
+    """Example: Load a CSV dataset from Kaggle."""
+    # Initialize loader (replace with your credentials)
+    loader = KaggleDataLoader("your_username", "your_api_key")
+    
+    # Download dataset (example: COVID-19 dataset)
+    dataset_path = loader.download_dataset("gpreda/covid-world-vaccination-progress")
+    
+    # Load CSV data
+    csv_file = os.path.join(dataset_path, "country_vaccinations.csv")
+    documents = loader.load_csv_dataset(
+        csv_file,
+        text_columns=["country", "vaccines", "source_name"],
+        chunk_size=100
+    )
+    
+    return documents
+
+def load_kaggle_json_example():
+    """Example: Load a JSON dataset from Kaggle."""
+    loader = KaggleDataLoader("your_username", "your_api_key")
+    
+    # Download dataset (example: news articles)
+    dataset_path = loader.download_dataset("rmisra/news-category-dataset")
+    
+    # Load JSON data
+    json_file = os.path.join(dataset_path, "News_Category_Dataset_v3.json")
+    documents = loader.load_json_dataset(
+        json_file,
+        text_field="headline",
+        metadata_fields=["category", "date"]
+    )
+    
+    return documents