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