{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.14","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"tpu1vmV38","dataSources":[{"sourceId":29705,"databundleVersionId":2662435,"sourceType":"competition"}],"dockerImageVersionId":30761,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"1. Setup and Imports\nBegin by importing the necessary libraries.","metadata":{}},{"cell_type":"code","source":"!pip install datatable\n!pip install pillow\n!pip install torch\n!pip install sentence-transformers==2.2.2","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-04T20:49:29.250460Z","iopub.execute_input":"2024-11-04T20:49:29.250856Z","iopub.status.idle":"2024-11-04T20:49:58.102979Z","shell.execute_reply.started":"2024-11-04T20:49:29.250821Z","shell.execute_reply":"2024-11-04T20:49:58.101885Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport os\nimport gzip\nimport base64\nfrom PIL import Image\nimport io\nimport re\n\n# For modeling\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\nimport torch.nn as nn\n\n# For text embeddings\nfrom sentence_transformers import SentenceTransformer\n\n# For similarity search\nfrom sklearn.metrics.pairwise import cosine_similarity\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-11-04T20:49:58.104836Z","iopub.execute_input":"2024-11-04T20:49:58.105120Z","iopub.status.idle":"2024-11-04T20:50:22.784876Z","shell.execute_reply.started":"2024-11-04T20:49:58.105091Z","shell.execute_reply":"2024-11-04T20:50:22.783926Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 2. Load a Subset of the Data\n2.1 Load Training Data\nSince the data is large, we'll read only a portion of the training data, focusing on English entries.","metadata":{}},{"cell_type":"code","source":"# Function to read a chunk of the TSV file\ndef read_tsv_in_chunks(file_path, chunk_size=10000, max_rows=50000):\n    chunks = []\n    total_rows = 0\n    for chunk in pd.read_csv(\n        file_path,\n        sep='\\t',\n        chunksize=chunk_size,\n        engine='python',\n        on_bad_lines='skip',\n        encoding='utf-8'\n    ):\n        # Filter English entries\n        english_chunk = chunk[chunk['language'] == 'en']\n        chunks.append(english_chunk)\n        total_rows += len(english_chunk)\n        # Break after collecting enough data\n        if total_rows >= max_rows:\n            break\n    return pd.concat(chunks, ignore_index=True)\n\n\n# Load a subset of the training data\ntrain_file = '/kaggle/input/wikipedia-image-caption/train-00000-of-00005.tsv'\ntrain_df = read_tsv_in_chunks(train_file)","metadata":{"execution":{"iopub.status.busy":"2024-11-04T20:50:22.786041Z","iopub.execute_input":"2024-11-04T20:50:22.786501Z","iopub.status.idle":"2024-11-04T20:50:39.055159Z","shell.execute_reply.started":"2024-11-04T20:50:22.786453Z","shell.execute_reply":"2024-11-04T20:50:39.054206Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-04T20:50:39.057144Z","iopub.execute_input":"2024-11-04T20:50:39.057452Z","iopub.status.idle":"2024-11-04T20:50:39.081285Z","shell.execute_reply.started":"2024-11-04T20:50:39.057423Z","shell.execute_reply":"2024-11-04T20:50:39.080663Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 2.2 Load Test Data\nSimilarly, load a small subset of the test data.","metadata":{}},{"cell_type":"code","source":"test_df = pd.read_csv('/kaggle/input/wikipedia-image-caption/test.tsv', sep='\\t')\n# For quick testing, select a small sample\ntest_df = test_df.sample(n=100, random_state=42).reset_index(drop=True)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-04T20:50:39.083544Z","iopub.execute_input":"2024-11-04T20:50:39.083837Z","iopub.status.idle":"2024-11-04T20:50:39.281919Z","shell.execute_reply.started":"2024-11-04T20:50:39.083807Z","shell.execute_reply":"2024-11-04T20:50:39.281105Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 2.3 Load Test Captions List","metadata":{}},{"cell_type":"code","source":"test_captions_df = pd.read_csv('/kaggle/input/wikipedia-image-caption/test_caption_list.csv')\n# For quick execution, you can sample or limit the captions\n# Here, we only consider English captions\ntest_captions_df = test_captions_df[test_captions_df['caption_title_and_reference_description'].apply(lambda x: x.isascii())]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-04T20:50:39.283005Z","iopub.execute_input":"2024-11-04T20:50:39.283271Z","iopub.status.idle":"2024-11-04T20:50:39.575606Z","shell.execute_reply.started":"2024-11-04T20:50:39.283245Z","shell.execute_reply":"2024-11-04T20:50:39.574660Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 3. Preprocess the Data\n3.1 Prepare Training Captions\nExtract relevant text fields and combine them to form the caption.","metadata":{}},{"cell_type":"code","source":"# Combine relevant text fields to create captions\ntrain_df['caption'] = train_df['caption_reference_description'].fillna('') + ' ' + \\\n                      train_df['caption_alt_text_description'].fillna('') + ' ' + \\\n                      train_df['page_title'].fillna('')\n\n# Keep only English captions\ntrain_df = train_df[train_df['caption'].apply(lambda x: x.isascii())].reset_index(drop=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-04T20:50:39.576961Z","iopub.execute_input":"2024-11-04T20:50:39.577230Z","iopub.status.idle":"2024-11-04T20:50:39.745252Z","shell.execute_reply.started":"2024-11-04T20:50:39.577198Z","shell.execute_reply":"2024-11-04T20:50:39.744254Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 3.2 Load Image Embeddings\nSince the raw images are large, we'll use the provided ResNet embeddings.\n","metadata":{}},{"cell_type":"code","source":"# Define the generator function\ndef load_resnet_embeddings_generator(file_paths, num_rows=None):\n    total_rows = 0\n\n    for file_path in file_paths:\n        with open(file_path, 'r', encoding='utf-8') as f:\n            for line_num, line in enumerate(f, 1):\n                line = line.strip()\n                if not line:\n                    continue\n\n                # Print the first few lines to verify the format\n                # if line_num <= 5:\n                #     print(f\"Line {line_num}: {line}\")\n\n                # Split the line on the first tab to separate URL and embeddings\n                if '\\t' in line:\n                    image_url, embedding_str = line.split('\\t', 1)\n                else:\n                    print(f\"Line {line_num}: No tab character found, skipping line.\")\n                    continue\n\n                # Split the embedding string on commas\n                embedding_values_str = embedding_str.strip().split(',')\n\n                # Convert embedding values to floats\n                try:\n                    embedding_values = [float(val) for val in embedding_values_str if val]\n                except ValueError as e:\n                    print(f\"Line {line_num}: Could not convert embedding values to floats. Error: {e}\")\n                    continue\n\n                # Convert to numpy array\n                embedding = np.array(embedding_values, dtype='float32')\n\n                # Yield the data\n                yield {'image_url': image_url, 'embedding': embedding}\n\n                total_rows += 1\n                if num_rows and total_rows >= num_rows:\n                    return\n\n# Prepare the file paths\nembedding_files = [\n    '/kaggle/input/wikipedia-image-caption/image_data_test/resnet_embeddings/test_resnet_embeddings_part-00000.csv'\n]\n\n# Create the generator\nembedding_generator = load_resnet_embeddings_generator(embedding_files, num_rows=5000)\n\n# Initialize an empty list or dict to store embeddings if needed\nembeddings_dict = {}\n","metadata":{"execution":{"iopub.status.busy":"2024-11-04T21:03:25.445974Z","iopub.execute_input":"2024-11-04T21:03:25.446887Z","iopub.status.idle":"2024-11-04T21:03:30.223817Z","shell.execute_reply.started":"2024-11-04T21:03:25.446849Z","shell.execute_reply":"2024-11-04T21:03:30.223001Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"for embedding_data in embedding_generator:\n    image_url = embedding_data['image_url']\n    embedding = embedding_data['embedding']\n    print(f\"Image URL: {image_url}\")\n    print(f\"Embedding shape: {embedding.shape}\")\n    # Store the embedding in a dict for quick access\n    embeddings_dict[image_url] = embedding\n    \n# embeddings_dict","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-04T21:06:31.135381Z","iopub.execute_input":"2024-11-04T21:06:31.136247Z","iopub.status.idle":"2024-11-04T21:06:31.140711Z","shell.execute_reply.started":"2024-11-04T21:06:31.136204Z","shell.execute_reply":"2024-11-04T21:06:31.140033Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Dataset Class\nDefine a dataset class to handle the loading and processing of images and captions.","metadata":{}},{"cell_type":"code","source":"class WikipediaDataset(Dataset):\n    def __init__(self, img_features, captions, tokenizer, max_length):\n        self.img_features = img_features\n        self.captions = captions\n        self.tokenizer = tokenizer\n        self.max_length = max_length\n\n    def __getitem__(self, idx):\n        if idx >= len(self.img_features):\n            idx = idx % len(self.img_features)\n\n        img_feature_data = self.img_features[idx]\n        img_url = img_feature_data[0].split('\\t')[0]\n        numeric_features = img_feature_data[1:].astype(np.float32)\n\n        try:\n            response = requests.get(img_url)\n            img = Image.open(BytesIO(response.content)).convert('RGB')\n        except (UnidentifiedImageError, requests.exceptions.RequestException):\n            # Skip this sample by returning None\n            return None\n\n        transform = transforms.Compose([\n            transforms.Resize((224, 224)),\n            transforms.ToTensor(),\n        ])\n        img_tensor = transform(img)\n\n        caption = self.captions[idx]\n        tokens = self.tokenizer(caption, max_length=self.max_length, padding='max_length', truncation=True, return_tensors='pt')\n\n        return {\n            'img_feature': img_tensor,\n            'numeric_features': torch.tensor(numeric_features, dtype=torch.float),\n            'ids': tokens['input_ids'].squeeze(0),\n        }\n    \n    def __len__(self):\n        return len(self.img_features)\n","metadata":{"execution":{"iopub.status.busy":"2024-11-04T20:50:42.010601Z","iopub.status.idle":"2024-11-04T20:50:42.010927Z","shell.execute_reply.started":"2024-11-04T20:50:42.010766Z","shell.execute_reply":"2024-11-04T20:50:42.010782Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Model Definition\nDefine a simple model that combines image features with text embeddings to output a similarity score.","metadata":{}},{"cell_type":"code","source":"class ImageCaptionModel(nn.Module):\n    def __init__(self, image_embedding_dim, text_embedding_dim, hidden_dim):\n        super(ImageCaptionModel, self).__init__()\n        self.image_fc = nn.Linear(image_embedding_dim, hidden_dim)\n        self.text_fc = nn.Linear(text_embedding_dim, hidden_dim)\n        self.relu = nn.ReLU()\n        self.output = nn.Linear(hidden_dim, 1)\n\n    def forward(self, img_feature, text_feature):\n        img_out = self.relu(self.image_fc(img_feature))\n        text_out = self.relu(self.text_fc(text_feature))\n        combined = img_out * text_out  # Element-wise multiplication\n        output = self.output(combined)\n        return output\n","metadata":{"execution":{"iopub.status.busy":"2024-11-04T20:50:42.012526Z","iopub.status.idle":"2024-11-04T20:50:42.012848Z","shell.execute_reply.started":"2024-11-04T20:50:42.012687Z","shell.execute_reply":"2024-11-04T20:50:42.012702Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Initialize Tokenizer and Model","metadata":{}},{"cell_type":"code","source":"# Set up the tokenizer and model\ntokenizer = BertTokenizer.from_pretrained('bert-base-uncased')\nbert_model = BertModel.from_pretrained('bert-base-uncased').to(device)\n\n# Initialize the model\nimage_embedding_dim = 2048\ntext_embedding_dim = 768\nhidden_dim = 512\n\nmodel = ImageCaptionModel(image_embedding_dim, text_embedding_dim, hidden_dim).to(device)\n\n","metadata":{"execution":{"iopub.status.busy":"2024-11-04T20:50:42.013763Z","iopub.status.idle":"2024-11-04T20:50:42.014083Z","shell.execute_reply.started":"2024-11-04T20:50:42.013908Z","shell.execute_reply":"2024-11-04T20:50:42.013927Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Initialize Optimizer","metadata":{}},{"cell_type":"code","source":"# Set up optimizer and loss function\noptimizer = optim.Adam(model.parameters(), lr=0.001)\ncriterion = nn.BCEWithLogitsLoss()","metadata":{"execution":{"iopub.status.busy":"2024-11-04T20:50:42.015696Z","iopub.status.idle":"2024-11-04T20:50:42.016060Z","shell.execute_reply.started":"2024-11-04T20:50:42.015877Z","shell.execute_reply":"2024-11-04T20:50:42.015896Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Prepare Training DataLoader","metadata":{}},{"cell_type":"code","source":"# Convert to pandas DataFrame if not already one\nif not isinstance(test_resnet_embeddings, pd.DataFrame):\n    test_resnet_embeddings = test_resnet_embeddings.to_pandas()\n\n# Set up the dataset and dataloader for training\nmax_length = 64  # Max length for the tokenized text\n\n# Prepare the dataset using the embeddings and captions\ntrain_dataset = WikipediaDataset(test_resnet_embeddings.values, test_captions_list, tokenizer, max_length)\n\n# Collate function to filter out None values\ndef collate_fn(batch):\n    batch = list(filter(lambda x: x is not None, batch))\n    if len(batch) == 0:\n        return None  # Return None for empty batches\n    return default_collate(batch)\n\n# Update DataLoader with the custom collate_fn\ntrain_loader = DataLoader(train_dataset, batch_size=8, shuffle=True, collate_fn=collate_fn)\n\nfor i, batch in enumerate(train_loader):\n    if batch is None:\n        continue\n    print(f\"Batch {i}: {batch['img_feature'].shape}, {batch['ids'].shape}\")","metadata":{"execution":{"iopub.status.busy":"2024-11-04T20:50:42.016989Z","iopub.status.idle":"2024-11-04T20:50:42.017293Z","shell.execute_reply.started":"2024-11-04T20:50:42.017140Z","shell.execute_reply":"2024-11-04T20:50:42.017155Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Training Loop with Memory Clearing and Smaller Batches","metadata":{}},{"cell_type":"code","source":"def train_model(model, dataloader, optimizer, criterion, epochs=5):\n    model.train()\n\n    for epoch in range(epochs):\n        total_loss = 0\n        for batch in dataloader:\n            img_features = batch['img_feature'].to(device)\n            ids = batch['ids'].to(device)\n            mask = batch['mask'].to(device)\n\n            with torch.no_grad():  # Disable autograd to save memory\n                text_embeddings = bert_model(input_ids=ids, attention_mask=mask)[0][:, 0, :]\n\n            optimizer.zero_grad()\n            outputs = model(img_features, text_embeddings)\n            labels = torch.ones(len(outputs)).to(device)\n            loss = criterion(outputs.view(-1), labels)\n\n            loss.backward()\n            optimizer.step()\n\n            # Clear memory after each batch\n            total_loss += loss.item()\n            clear_memory()\n\n        print(f'Epoch {epoch + 1}/{epochs}, Loss: {total_loss / len(dataloader)}')\n        torch.save(model.state_dict(), f'/kaggle/working/model_checkpoint_epoch_{epoch+1}.pth')  # Save after each epoch\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-04T20:50:42.018032Z","iopub.status.idle":"2024-11-04T20:50:42.018328Z","shell.execute_reply.started":"2024-11-04T20:50:42.018175Z","shell.execute_reply":"2024-11-04T20:50:42.018191Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Train the Model","metadata":{}},{"cell_type":"code","source":"# Train the model\ntrain_model(model, train_loader, optimizer, criterion, epochs=5)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-04T20:50:42.019195Z","iopub.status.idle":"2024-11-04T20:50:42.019514Z","shell.execute_reply.started":"2024-11-04T20:50:42.019352Z","shell.execute_reply":"2024-11-04T20:50:42.019367Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Inference and Submission\nGenerate predictions for the test set and format them for submission.","metadata":{}},{"cell_type":"code","source":"# Prepare for inference\nmodel.eval()\nsubmission = []\n\nwith torch.no_grad():  # No need for gradients during inference\n    for i, batch in enumerate(train_loader):\n        img_features = batch['img_feature'].to(device)\n        ids = batch['ids'].to(device)\n        mask = batch['mask'].to(device)\n\n        # Perform inference\n        text_embeddings = bert_model(input_ids=ids, attention_mask=mask)[0][:, 0, :]\n        outputs = model(img_features, text_embeddings).cpu().numpy()\n        sorted_indices = np.argsort(outputs, axis=0)[::-1]\n        top_captions = [test_captions_list[idx] for idx in sorted_indices[:5]]\n\n        for caption in top_captions:\n            submission.append({'id': i, 'caption_title_and_reference_description': caption})\n\n# Save submission file\nsubmission_df = pd.DataFrame(submission)\nsubmission_df.to_csv('/kaggle/working/submission.csv', index=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-04T20:50:42.021032Z","iopub.status.idle":"2024-11-04T20:50:42.021330Z","shell.execute_reply.started":"2024-11-04T20:50:42.021182Z","shell.execute_reply":"2024-11-04T20:50:42.021196Z"}},"outputs":[],"execution_count":null}]}