Spaces:
Runtime error
Runtime error
| """ | |
| ML Pipeline Data Loader for ScamDetect AI | |
| Handles the 50,000+ research-level text, URL, image, and video datasets. | |
| Provides PyTorch Dataset and DataLoader abstractions. | |
| """ | |
| import os | |
| import pandas as pd | |
| try: | |
| import torch | |
| from torch.utils.data import Dataset, DataLoader | |
| from transformers import AutoTokenizer | |
| TORCH_AVAILABLE = True | |
| except ImportError: | |
| TORCH_AVAILABLE = False | |
| print("Warning: PyTorch not installed. ML DataLoader will operate in pandas-only mode.") | |
| DATA_DIR = os.path.join(os.path.dirname(__file__), "data") | |
| class TextScamDataset(Dataset if TORCH_AVAILABLE else object): | |
| def __init__(self, csv_file=None, tokenizer_name="bert-base-multilingual-cased", max_length=128): | |
| self.csv_file = csv_file or os.path.join(DATA_DIR, "text_dataset.csv") | |
| self.data = pd.read_csv(self.csv_file) | |
| # Mapping labels | |
| self.label_map = {"safe": 0, "scam": 1} | |
| self.max_length = max_length | |
| if TORCH_AVAILABLE: | |
| self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name) | |
| def __len__(self): | |
| return len(self.data) | |
| def __getitem__(self, idx): | |
| row = self.data.iloc[idx] | |
| text = str(row['text']) | |
| label = self.label_map.get(row['label'], 1) | |
| if not TORCH_AVAILABLE: | |
| return {"text": text, "label": label, "category": row['category']} | |
| encoding = self.tokenizer( | |
| text, | |
| add_special_tokens=True, | |
| max_length=self.max_length, | |
| return_token_type_ids=False, | |
| padding='max_length', | |
| truncation=True, | |
| return_attention_mask=True, | |
| return_tensors='pt', | |
| ) | |
| return { | |
| 'text': text, | |
| 'input_ids': encoding['input_ids'].flatten(), | |
| 'attention_mask': encoding['attention_mask'].flatten(), | |
| 'targets': torch.tensor(label, dtype=torch.long), | |
| 'category': row['category'] | |
| } | |
| class UrlPhishingDataset(Dataset if TORCH_AVAILABLE else object): | |
| def __init__(self, csv_file=None): | |
| self.csv_file = csv_file or os.path.join(DATA_DIR, "url_dataset.csv") | |
| self.data = pd.read_csv(self.csv_file) | |
| self.label_map = {"safe": 0, "scam": 1} | |
| def __len__(self): | |
| return len(self.data) | |
| def __getitem__(self, idx): | |
| row = self.data.iloc[idx] | |
| url = str(row['url']) | |
| label = self.label_map.get(row['label'], 1) | |
| item = {"url": url, "label": label} | |
| if TORCH_AVAILABLE: | |
| item["targets"] = torch.tensor(label, dtype=torch.long) | |
| return item | |
| def get_data_loaders(batch_size=32): | |
| """Returns PyTorch DataLoaders for train/test splits.""" | |
| if not TORCH_AVAILABLE: | |
| raise ImportError("PyTorch required for DataLoader generation") | |
| text_ds = TextScamDataset() | |
| url_ds = UrlPhishingDataset() | |
| # In a real scenario, split train/test here using torch.utils.data.random_split | |
| text_loader = DataLoader(text_ds, batch_size=batch_size, shuffle=True) | |
| url_loader = DataLoader(url_ds, batch_size=batch_size, shuffle=True) | |
| return {"text": text_loader, "url": url_loader} | |
| if __name__ == "__main__": | |
| print("Loading datasets...") | |
| text_ds = TextScamDataset() | |
| url_ds = UrlPhishingDataset() | |
| print(f"Loaded Text Dataset: {len(text_ds)} samples") | |
| print(f"Loaded URL Dataset: {len(url_ds)} samples") | |
| if len(text_ds) > 0: | |
| print("\nSample Text Record:") | |
| print(text_ds[0]) | |