Add test dataset and example inference/evaluation for LSTM model

.gitignore

@@ -7,4 +7,5 @@ __pycache__/
 *.db
 metrics.txt
 predictions.txt
-*.pth
+*.pth
+news-category-dataset/

README.md

@@ -1,32 +1,9 @@
-# DocBERT - Improved Document Classification with BERT
-
-This repository contains an improved implementation of BERT for document classification, combining techniques from [jesse-tong/docbert](https://github.com/jesse-tong/docbert) and [castorini/hedwig](https://github.com/castorini/hedwig).
-
-## Key Improvements
-
-1. **Advanced Regularization Techniques**:
-   - Dropout in multiple layers
-   - Layer normalization
-   - Gradient clipping
-   - Weight decay optimization
-
-2. **Training Stability Enhancements**:
-   - Learning rate scheduling with ReduceLROnPlateau
-   - Gradient accumulation for effective larger batch sizes
-   - Label smoothing to improve generalization
-   - Early stopping based on validation F1 score
-
-3. **Architectural Changes**:
-   - Better BERT pooling strategies
-   - More robust tokenization with attention masks
-   - Configurable hyperparameters for different document types
-
 ## Installation
 
 ```bash
 # Clone the repository
-git clone https://github.com/yourusername/docbert-improved.git
-cd docbert-improved
+git clone https://github.com/jesse-tong/docbert2.git
+cd docbert2
 
 # Install dependencies
 pip install -r requirements.txt

dataset_lstm.py

@@ -99,10 +99,17 @@ class LSTMDataset(Dataset):
             'attention_mask': encoding['attention_mask'],
             'label': torch.tensor(label, dtype=torch.long)
         }
+    
+    def get_text_(self, idx):
+        """Get original text for a given index"""
+        return {
+            'text': self.texts[idx],
+            'label': self.labels[idx]
+        }
 
 def prepare_lstm_data(data_path, text_col='text', label_col='label', 
                      max_vocab_size=30000, max_seq_length=512,
-                     val_split=0.1, test_split=0.1, batch_size=32, seed=42):
+                     val_split=0.1, test_split=0.1, batch_size=32, seed=42, return_datasets=False):
     """
     Load data and prepare for LSTM model
     """
@@ -154,10 +161,27 @@ def prepare_lstm_data(data_path, text_col='text', label_col='label',
     train_dataset = LSTMDataset(train_texts, train_labels, tokenizer)
     val_dataset = LSTMDataset(val_texts, val_labels, tokenizer)
     test_dataset = LSTMDataset(test_texts, test_labels, tokenizer)
+
+    if return_datasets:
+        return train_dataset, val_dataset, test_dataset, tokenizer.vocab_size
     
     # Create data loaders
-    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
-    val_loader = DataLoader(val_dataset, batch_size=batch_size)
-    test_loader = DataLoader(test_dataset, batch_size=batch_size)
+    if len(train_dataset.texts) == 0:
+        logger.warning("Training dataset is empty. Please check your data.")
+        train_loader = None
+    else:
+        train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
+    
+    if len(val_dataset.texts) == 0:
+        logger.warning("Validation dataset is empty. Please check your data.")
+        val_loader = None
+    else:
+        val_loader = DataLoader(val_dataset, batch_size=batch_size)
+    
+    if len(test_dataset.texts) == 0:
+        logger.warning("Test dataset is empty. Please check your data.")
+        test_loader = None
+    else:
+        test_loader = DataLoader(test_dataset, batch_size=batch_size)
     
     return train_loader, val_loader, test_loader, tokenizer.vocab_size

example_uses.md

@@ -5,12 +5,17 @@
 ```
 python train.py --data_path train.csv --label_column "Class Index" --text_column "Description" --epochs 4 --num_classes 4
 ```
-- Inference with BERT model (train.csv is ag_news dataset with 4 classes)
+- Inference with BERT model (test_data.csv is test dataset with 4 classes like ag_news)
 ```
-python .\inference_example.py --model_path "./bert_base_uncased/best_model.pth" --num_classes 4 --class_names "World" "Sports" "Business" "Science" --text_column "Description" --label_column "Class Index" --data_path "./train.csv" --inference_batch_limit 10
+python .\inference_example.py --model_path "./bert_base_uncased/best_model.pth" --num_classes 4 --class_names "World" "Sports" "Business" "Science" --text_column "Description" --label_column "Class Index" --data_path "./test_data.csv" --inference_batch_limit 10
 ```
 
-- Train LSTM model from BERT model using distillation
+- Train LSTM model from BERT model using distillation (train dataset should be the same as distillation training dataset)
 ```
 python .\distill_bert_to_lstm.py --bert_model bert-base-uncased --bert_model_path "./bert_base_uncased/best_model.pth" --output_dir "./docbert_lstm" --batch_size 32 --epochs 10 --data_path "./train.csv" --text_column "Description" --label_column "Class Index" --num_classes 4
+```
+
+- Inference with distilled LSTM model (test_data.csv is test dataset with 4 classes like ag_news)
+```
+python .\inference_lstm.py --model_path "./docbert_lstm/distilled_lstm_model.pth" --num_classes 4 --class_names "World" "Sports" "Business" "Science" --text_column "Description" --label_column "Class Index" --data_path "./test_data.csv" --inference_batch_limit 10
 ```

inference_lstm.py

@@ -0,0 +1,100 @@
+from dataset_lstm import prepare_lstm_data, LSTMTokenizer, LSTMDataset
+from models.lstm_model import DocumentBiLSTM
+import sklearn, torch
+import numpy as np
+import argparse
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Document Classification with LSTM")
+    parser.add_argument("--data_path", type=str, required=True, help="Path to the dataset")
+    parser.add_argument("--model_path", type=str, required=True, help="Path to the trained model")
+    parser.add_argument("--max_seq_length", type=int, default=512, help="Maximum sequence length for LSTM")
+    parser.add_argument("--batch_size", type=int, default=32, help="Batch size for training and evaluation")
+    parser.add_argument("--num_classes", type=int, required=True, help="Number of classes for classification")
+    parser.add_argument("--text_column", type=str, default="text", help="Column name for text data")
+    parser.add_argument("--label_column", type=str, default="label", help="Column name for labels")
+    parser.add_argument("--class_names", type=str, nargs='+', required=True, help="List of class names for classification")
+    parser.add_argument("--inference_batch_limit", type=int, default=-1, help="Limit for inference batch counts")
+    parser.add_argument("--print_predictions", type=bool, default=False, help="Print predictions to console")
+    args = parser.parse_args()
+
+    class_names = args.class_names
+
+    # Set device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # Prepare data
+    tokenizer = LSTMTokenizer(max_length=args.max_seq_length)
+    test_dataset, _, _, vocab_size = prepare_lstm_data(args.data_path,
+                                    text_col=args.text_column,
+                                    label_col=args.label_column,
+                                    tokenizer=tokenizer,
+                                    batch_size=args.batch_size,
+                                    max_length=args.max_seq_length, val_split=0.0, test_split=0.0, return_datasets=True)
+
+    test_loader, _, _, vocab_size = prepare_lstm_data(args.data_path,
+                                  text_col=args.text_column,
+                                  label_col=args.label_column,
+                                  tokenizer=tokenizer,
+                                  batch_size=args.batch_size,
+                                  max_length=args.max_seq_length, val_split=0.0, test_split=0.0)
+    
+    # Load model
+    model = DocumentBiLSTM(input_dim=tokenizer.vocab_size,
+                           embedding_dim=128,
+                           hidden_dim=64,
+                           output_dim=args.num_classes)
+    
+    model.load_state_dict(torch.load(args.model_path))
+    model = model.to(device)
+
+    all_labels = np.array([], dtype=int)
+    all_predictions = np.array([], dtype=int)
+    
+    # Inference
+    batch_count = 0
+    with torch.no_grad():
+        for batch in test_loader:
+            input_ids = batch['input_ids'].to(device)
+            labels = batch['label'].to(device)
+            all_labels = np.append(all_labels, labels.cpu().numpy())
+
+            outputs = model(input_ids)
+            predictions = torch.argmax(outputs, dim=1)
+            all_predictions = np.append(all_predictions, predictions.cpu().numpy())
+
+            if args.print_predictions:
+                for i in range(len(predictions)):
+                    print(f"Text: {test_dataset.get_text_(batch_count * args.batch_size + i)}, Prediction: {class_names[predictions[i]]}, True Label: {class_names[labels[i]]}")
+                
+            if args.inference_batch_limit > 0 and batch_count >= args.inference_batch_limit:
+                break
+
+            batch_count += 1
+
+    # Print classification report
+    # Calculate accuracy, F1 score, recall, and precision
+    accuracy = sklearn.metrics.accuracy_score(all_labels, all_predictions)
+    f1 = sklearn.metrics.f1_score(all_labels, all_predictions, average='weighted')
+    precision = sklearn.metrics.precision_score(all_labels, all_predictions, average='weighted')
+    recall = sklearn.metrics.recall_score(all_labels, all_predictions, average='weighted')
+
+    print(f"Accuracy: {accuracy}")
+    print(f"F1 Score: {f1}")
+    print(f"Precision: {precision}")
+    print(f"Recall: {recall}")
+
+    with open("predictions_lstm.txt", "w") as f:
+        for i in range(len(all_labels)):
+            idx = int(i)
+            f.write(f"Text: {test_dataset.get_text_(idx)}\n")
+            f.write(f"True Label: {all_labels[idx]}, Predicted Label: {all_predictions[idx]}\n")
+            f.write(f"Predicted Class: {class_names[all_predictions[idx]] if len(class_names) > all_predictions[idx] else "Unknown"}, True Class: {class_names[all_labels[idx]] if len(class_names) > all_predictions[idx] else "Unknown"}\n")
+            f.write("\n")
+
+    with open("metrics_lstm.txt", "w") as f:
+        f.write(f"Accuracy: {accuracy}\n")
+        f.write(f"F1 Score: {f1}\n")
+        f.write(f"Precision: {precision}\n")
+        f.write(f"Recall: {recall}\n")
+            

test_data_creation.ipynb

@@ -0,0 +1,476 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "c:\\Users\\admin\\Documents\\DocBERT2\\news-category-dataset\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Cloning into 'news-category-dataset'...\n",
+      "c:\\Users\\admin\\AppData\\Local\\Programs\\Python\\Python312\\Lib\\site-packages\\IPython\\core\\magics\\osm.py:417: UserWarning: This is now an optional IPython functionality, setting dhist requires you to install the `pickleshare` library.\n",
+      "  self.shell.db['dhist'] = compress_dhist(dhist)[-100:]\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Dataset is from Heegyu Kim, available at https://huggingface.co/datasets/heegyu/news-category-dataset, under CC-BY 4.0 license.\n",
+    "# The dataset is a collection of news articles from various categories from Huffington Post.\n",
+    "!git clone https://huggingface.co/datasets/heegyu/news-category-dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pandas as pd\n",
+    "import numpy as np"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>link</th>\n",
+       "      <th>headline</th>\n",
+       "      <th>category</th>\n",
+       "      <th>short_description</th>\n",
+       "      <th>authors</th>\n",
+       "      <th>date</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>https://www.huffpost.com/entry/covid-boosters-...</td>\n",
+       "      <td>Over 4 Million Americans Roll Up Sleeves For O...</td>\n",
+       "      <td>U.S. NEWS</td>\n",
+       "      <td>Health experts said it is too early to predict...</td>\n",
+       "      <td>Carla K. Johnson, AP</td>\n",
+       "      <td>2022-09-23</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>https://www.huffpost.com/entry/american-airlin...</td>\n",
+       "      <td>American Airlines Flyer Charged, Banned For Li...</td>\n",
+       "      <td>U.S. NEWS</td>\n",
+       "      <td>He was subdued by passengers and crew when he ...</td>\n",
+       "      <td>Mary Papenfuss</td>\n",
+       "      <td>2022-09-23</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>https://www.huffpost.com/entry/funniest-tweets...</td>\n",
+       "      <td>23 Of The Funniest Tweets About Cats And Dogs ...</td>\n",
+       "      <td>COMEDY</td>\n",
+       "      <td>\"Until you have a dog you don't understand wha...</td>\n",
+       "      <td>Elyse Wanshel</td>\n",
+       "      <td>2022-09-23</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>https://www.huffpost.com/entry/funniest-parent...</td>\n",
+       "      <td>The Funniest Tweets From Parents This Week (Se...</td>\n",
+       "      <td>PARENTING</td>\n",
+       "      <td>\"Accidentally put grown-up toothpaste on my to...</td>\n",
+       "      <td>Caroline Bologna</td>\n",
+       "      <td>2022-09-23</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>https://www.huffpost.com/entry/amy-cooper-lose...</td>\n",
+       "      <td>Woman Who Called Cops On Black Bird-Watcher Lo...</td>\n",
+       "      <td>U.S. NEWS</td>\n",
+       "      <td>Amy Cooper accused investment firm Franklin Te...</td>\n",
+       "      <td>Nina Golgowski</td>\n",
+       "      <td>2022-09-22</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "                                                link  \\\n",
+       "0  https://www.huffpost.com/entry/covid-boosters-...   \n",
+       "1  https://www.huffpost.com/entry/american-airlin...   \n",
+       "2  https://www.huffpost.com/entry/funniest-tweets...   \n",
+       "3  https://www.huffpost.com/entry/funniest-parent...   \n",
+       "4  https://www.huffpost.com/entry/amy-cooper-lose...   \n",
+       "\n",
+       "                                            headline   category  \\\n",
+       "0  Over 4 Million Americans Roll Up Sleeves For O...  U.S. NEWS   \n",
+       "1  American Airlines Flyer Charged, Banned For Li...  U.S. NEWS   \n",
+       "2  23 Of The Funniest Tweets About Cats And Dogs ...     COMEDY   \n",
+       "3  The Funniest Tweets From Parents This Week (Se...  PARENTING   \n",
+       "4  Woman Who Called Cops On Black Bird-Watcher Lo...  U.S. NEWS   \n",
+       "\n",
+       "                                   short_description               authors  \\\n",
+       "0  Health experts said it is too early to predict...  Carla K. Johnson, AP   \n",
+       "1  He was subdued by passengers and crew when he ...        Mary Papenfuss   \n",
+       "2  \"Until you have a dog you don't understand wha...         Elyse Wanshel   \n",
+       "3  \"Accidentally put grown-up toothpaste on my to...      Caroline Bologna   \n",
+       "4  Amy Cooper accused investment firm Franklin Te...        Nina Golgowski   \n",
+       "\n",
+       "        date  \n",
+       "0 2022-09-23  \n",
+       "1 2022-09-23  \n",
+       "2 2022-09-23  \n",
+       "3 2022-09-23  \n",
+       "4 2022-09-22  "
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "df = pd.read_json(\"./news-category-dataset/data.json\", lines=True)\n",
+    "df.head(5)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Unique categories in the dataset:  ['U.S. NEWS' 'COMEDY' 'PARENTING' 'WORLD NEWS' 'CULTURE & ARTS' 'TECH'\n",
+      " 'SPORTS' 'ENTERTAINMENT' 'POLITICS' 'WEIRD NEWS' 'ENVIRONMENT'\n",
+      " 'EDUCATION' 'CRIME' 'SCIENCE' 'WELLNESS' 'BUSINESS' 'STYLE & BEAUTY'\n",
+      " 'FOOD & DRINK' 'MEDIA' 'QUEER VOICES' 'HOME & LIVING' 'WOMEN'\n",
+      " 'BLACK VOICES' 'TRAVEL' 'MONEY' 'RELIGION' 'LATINO VOICES' 'IMPACT'\n",
+      " 'WEDDINGS' 'COLLEGE' 'PARENTS' 'ARTS & CULTURE' 'STYLE' 'GREEN' 'TASTE'\n",
+      " 'HEALTHY LIVING' 'THE WORLDPOST' 'GOOD NEWS' 'WORLDPOST' 'FIFTY' 'ARTS'\n",
+      " 'DIVORCE']\n"
+     ]
+    }
+   ],
+   "source": [
+    "unique_categories = df['category'].unique()\n",
+    "print(\"Unique categories in the dataset: \", unique_categories)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Number of WORLD category articles:  9542\n"
+     ]
+    }
+   ],
+   "source": [
+    "world_category = df[(df['category'] == 'WORLD') | (df['category'] == 'WORLD NEWS') | (df['category'] == 'WORLDPOST') | (df['category'] == 'THE WORLDPOST')]\n",
+    "print(\"Number of WORLD category articles: \", len(world_category))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Number of SPORTS category articles:  5077\n"
+     ]
+    }
+   ],
+   "source": [
+    "sports_category = df[(df['category'] == 'SPORTS') | (df['category'] == 'SPORT')]\n",
+    "print(\"Number of SPORTS category articles: \", len(sports_category))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Number of BUSINESS category articles:  5992\n"
+     ]
+    }
+   ],
+   "source": [
+    "business_category = df[(df['category'] == 'BUSINESS')]\n",
+    "print(\"Number of BUSINESS category articles: \", len(business_category))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Number of SCIENCE category articles:  4310\n"
+     ]
+    }
+   ],
+   "source": [
+    "science_category = df[(df['category'] == 'SCIENCE') | (df['category'] == 'SCIENCE & TECH') | (df['category'] == 'SCIENCE & TECH') | (df['category'] == 'TECH')]\n",
+    "print(\"Number of SCIENCE category articles: \", len(science_category))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "C:\\Users\\admin\\AppData\\Local\\Temp\\ipykernel_6844\\3064268492.py:1: SettingWithCopyWarning: \n",
+      "A value is trying to be set on a copy of a slice from a DataFrame.\n",
+      "Try using .loc[row_indexer,col_indexer] = value instead\n",
+      "\n",
+      "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n",
+      "  science_category['Description'] = science_category['headline'] + \"\\n\\n\" + science_category['short_description']\n",
+      "C:\\Users\\admin\\AppData\\Local\\Temp\\ipykernel_6844\\3064268492.py:2: SettingWithCopyWarning: \n",
+      "A value is trying to be set on a copy of a slice from a DataFrame.\n",
+      "Try using .loc[row_indexer,col_indexer] = value instead\n",
+      "\n",
+      "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n",
+      "  science_category['Class Index'] = 3\n",
+      "C:\\Users\\admin\\AppData\\Local\\Temp\\ipykernel_6844\\3064268492.py:4: SettingWithCopyWarning: \n",
+      "A value is trying to be set on a copy of a slice from a DataFrame.\n",
+      "Try using .loc[row_indexer,col_indexer] = value instead\n",
+      "\n",
+      "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n",
+      "  business_category['Description'] = business_category['headline'] + \"\\n\\n\" + business_category['short_description']\n",
+      "C:\\Users\\admin\\AppData\\Local\\Temp\\ipykernel_6844\\3064268492.py:5: SettingWithCopyWarning: \n",
+      "A value is trying to be set on a copy of a slice from a DataFrame.\n",
+      "Try using .loc[row_indexer,col_indexer] = value instead\n",
+      "\n",
+      "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n",
+      "  business_category['Class Index'] = 2\n",
+      "C:\\Users\\admin\\AppData\\Local\\Temp\\ipykernel_6844\\3064268492.py:7: SettingWithCopyWarning: \n",
+      "A value is trying to be set on a copy of a slice from a DataFrame.\n",
+      "Try using .loc[row_indexer,col_indexer] = value instead\n",
+      "\n",
+      "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n",
+      "  sports_category['Description'] = sports_category['headline'] + \"\\n\\n\" + sports_category['short_description']\n",
+      "C:\\Users\\admin\\AppData\\Local\\Temp\\ipykernel_6844\\3064268492.py:8: SettingWithCopyWarning: \n",
+      "A value is trying to be set on a copy of a slice from a DataFrame.\n",
+      "Try using .loc[row_indexer,col_indexer] = value instead\n",
+      "\n",
+      "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n",
+      "  sports_category['Class Index'] = 1\n",
+      "C:\\Users\\admin\\AppData\\Local\\Temp\\ipykernel_6844\\3064268492.py:10: SettingWithCopyWarning: \n",
+      "A value is trying to be set on a copy of a slice from a DataFrame.\n",
+      "Try using .loc[row_indexer,col_indexer] = value instead\n",
+      "\n",
+      "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n",
+      "  world_category['Description'] = world_category['headline'] + \"\\n\\n\" + world_category['short_description']\n",
+      "C:\\Users\\admin\\AppData\\Local\\Temp\\ipykernel_6844\\3064268492.py:11: SettingWithCopyWarning: \n",
+      "A value is trying to be set on a copy of a slice from a DataFrame.\n",
+      "Try using .loc[row_indexer,col_indexer] = value instead\n",
+      "\n",
+      "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n",
+      "  world_category['Class Index'] = 0\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>Description</th>\n",
+       "      <th>Class Index</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>Trump’s Saudi Trip Should Be About Ending The ...</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>Russia Vows To Expand 'Black List' Of American...</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>It's Been A Long, Crazy Year Since Britain's S...</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>First Penis Transplants In U.S. Planned For Wo...</td>\n",
+       "      <td>3</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>U.S.-Backed Forces Prepare For ISIS To Use Che...</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>5</th>\n",
+       "      <td>The European Migrant Crisis Is A Nightmare. Cl...</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6</th>\n",
+       "      <td>A Bunch Of Strange Things Just Went Down At NB...</td>\n",
+       "      <td>1</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>7</th>\n",
+       "      <td>Minor Explosion Rattles IMF Headquarters In Pa...</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>8</th>\n",
+       "      <td>Shalane Flanagan Becomes First U.S. Woman To W...</td>\n",
+       "      <td>1</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>9</th>\n",
+       "      <td>10 Instagrams To Celebrate Crown Prince Hussei...</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "                                         Description  Class Index\n",
+       "0  Trump’s Saudi Trip Should Be About Ending The ...            0\n",
+       "1  Russia Vows To Expand 'Black List' Of American...            0\n",
+       "2  It's Been A Long, Crazy Year Since Britain's S...            0\n",
+       "3  First Penis Transplants In U.S. Planned For Wo...            3\n",
+       "4  U.S.-Backed Forces Prepare For ISIS To Use Che...            0\n",
+       "5  The European Migrant Crisis Is A Nightmare. Cl...            0\n",
+       "6  A Bunch Of Strange Things Just Went Down At NB...            1\n",
+       "7  Minor Explosion Rattles IMF Headquarters In Pa...            0\n",
+       "8  Shalane Flanagan Becomes First U.S. Woman To W...            1\n",
+       "9  10 Instagrams To Celebrate Crown Prince Hussei...            0"
+      ]
+     },
+     "execution_count": 18,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "science_category['Description'] = science_category['headline'] + \"\\n\\n\" + science_category['short_description']\n",
+    "science_category['Class Index'] = 3\n",
+    "\n",
+    "business_category['Description'] = business_category['headline'] + \"\\n\\n\" + business_category['short_description']\n",
+    "business_category['Class Index'] = 2\n",
+    "\n",
+    "sports_category['Description'] = sports_category['headline'] + \"\\n\\n\" + sports_category['short_description']\n",
+    "sports_category['Class Index'] = 1\n",
+    "\n",
+    "world_category['Description'] = world_category['headline'] + \"\\n\\n\" + world_category['short_description']\n",
+    "world_category['Class Index'] = 0\n",
+    "\n",
+    "science_category = science_category[['Description', 'Class Index']]\n",
+    "business_category = business_category[['Description', 'Class Index']]\n",
+    "sports_category = sports_category[['Description', 'Class Index']]\n",
+    "world_category = world_category[['Description', 'Class Index']]\n",
+    "\n",
+    "test_data_df = pd.concat([science_category, business_category, sports_category, world_category], ignore_index=True)\n",
+    "# Shuffle the DataFrame\n",
+    "test_data_df = test_data_df.sample(frac=1, random_state=42).reset_index(drop=True)\n",
+    "\n",
+    "test_data_df.head(10)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "test_data_df.to_csv(\"test_data.csv\", index=False)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.1"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}