graph-inspect/inspect-graph-comparison.ipynb: update

dataset-retrieval/retrieving-bluesky-datasets.ipynb

@@ -98,14 +98,16 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import re, random, time\n",
     "import pandas as pd\n",
     "import numpy as np\n",
-    "import re\n",
     "from tqdm import tqdm\n",
     "from more_itertools import chunked\n",
     "\n",
     "# Remove URLs, hashtags, mentions, emojis and whitespaces\n",
     "def clean_text(text):\n",
+    "    if not isinstance(text, str):\n",
+    "        return np.nan\n",
     "    text = re.sub(r'http\\S+|www\\S+|https\\S+', '', text, flags=re.MULTILINE)\n",
     "    text = re.sub(r'#\\S+', '', text)\n",
     "    text = re.sub(r'@\\S+', '', text)\n",
@@ -113,110 +115,163 @@
     "    text = re.sub(r'\\s+', ' ', text).strip()\n",
     "    return text\n",
     "\n",
-    "# Search for posts containing a specific keyword\n",
+    "# Retry helper\n",
-    "def get_posts_by_keyword(keyword, since, until, limit=300):\n",
+    "def call_with_retry(fn, *args, max_retries=5, base_delay=2, **kwargs):\n",
+    "    last_exc = None\n",
+    "    for attempt in range(1, max_retries + 1):\n",
     "        try:\n",
-    "        num_posts = 0\n",
+    "            return fn(*args, **kwargs)\n",
-    "        post_data = []\n",
+    "        except Exception as e:\n",
+    "            last_exc = e\n",
+    "            wait = base_delay * (2 ** (attempt - 1)) + random.uniform(0, 1)\n",
+    "            print(f' [retry {attempt}/{max_retries}] {type(e).__name__}: {e} - retrying in {wait:.1f}s')\n",
+    "            time.sleep(wait)\n",
+    "    raise last_exc\n",
     "\n",
-    "        progress_bar = tqdm(total=limit, desc=\"Fetching posts\")\n",
+    "# Post field extraction\n",
-    "\n",
+    "def extract_post(post):\n",
-    "        while num_posts < limit:\n",
+    "    try:\n",
+    "        text = clean_text(post.record.text) if hasattr(post.record, 'text') else None\n",
+    "        iso_date = post.record.created_at if hasattr(post.record, 'created_at') else None\n",
+    "        user = post.author.handle if hasattr(post.author, 'handle') else None\n",
+    "        post_id = post.uri if hasattr(post, 'uri') else None\n",
+    "        parent_id = (\n",
+    "            post.record.reply.parent.uri\n",
+    "            if hasattr(post.record, 'reply') and hasattr(post.record.reply, 'parent')\n",
+    "            else 'no_parent'\n",
+    "        )\n",
+    " \n",
+    "        if iso_date is None or post_id is None or user is None:\n",
+    "            return None\n",
     " \n",
-    "            search_results = client.app.bsky.feed.search_posts({'q': keyword, 'since': since, 'until': until, 'lang': 'en', 'limit': 100})\n",
-    "            posts = search_results.posts\n",
-    "\n",
-    "            for post in posts:\n",
-    "                        # Extract post details\n",
-    "                        text = clean_text(post.record.text) if hasattr(post.record, 'text') else np.nan\n",
-    "                        iso_date = post.record.created_at if hasattr(post.record, 'created_at') else np.nan\n",
-    "                        user = post.author.handle if hasattr(post.author, 'handle') else np.nan\n",
-    "\n",
-    "                        post_id = post.uri if hasattr(post, 'uri') else np.nan\n",
-    "                        parent_id = post.record.reply.parent.uri if hasattr(post.record, 'reply') and hasattr(post.record.reply, 'parent') else 'no_parent'\n",
-    "\n",
-    "                        # Format date-time\n",
-    "                        if iso_date != np.nan:\n",
     "        date_obj = datetime.fromisoformat(iso_date.replace('Z', '+00:00'))\n",
     "        date_str = date_obj.strftime('%Y-%m-%d %H:%M:%S')\n",
-    "                        else:\n",
+    " \n",
-    "                            date_str = np.nan\n",
+    "        return {\n",
-    "\n",
-    "                        if post_id != np.nan and user != np.nan:\n",
-    "                            post_data.append({\n",
     "            'date': date_str,\n",
-    "                                'text': text,\n",
+    "            'text': text if text is not None else np.nan,\n",
     "            'user': user,\n",
     "            'id': post_id,\n",
-    "                                'parentid': parent_id\n",
+    "            'parentid': parent_id,\n",
-    "                            })\n",
+    "        }, date_obj\n",
-    "                        else:\n",
-    "                            continue\n",
-    "                        \n",
-    "            until = iso_date\n",
-    "            num_posts += len(posts)\n",
-    "            progress_bar.update(len(posts))\n",
-    "            progress_bar.set_postfix_str(f'now at: {date_str}')\n",
-    "\n",
-    "            if date_obj < (start_timestamp + timedelta(hours=1)):\n",
-    "                break\n",
-    "\n",
-    "        progress_bar.close()\n",
-    "        df = pd.DataFrame(post_data)\n",
-    "        return df\n",
-    "\n",
     "    except Exception as e:\n",
-    "        print(f\"Error fetching posts: {e}\")\n",
+    "        print(f\"  Skipping malformed post: {e}\")\n",
+    "        return None\n",
+    "\n",
+    "# Cleanup\n",
+    "def _dedupe(rows):\n",
+    "    if not rows:\n",
     "        return pd.DataFrame()\n",
+    "    df = pd.DataFrame(rows)\n",
+    "    before = len(df)\n",
+    "    df = df.drop_duplicates(subset='id')\n",
+    "    df = df.dropna(subset=['text', 'id', 'user'])\n",
+    "    df = df[df['text'].astype(str).str.strip() != '']\n",
+    "    after = len(df)\n",
+    "    if before != after:\n",
+    "        print(f\"(Filtered {before - after} duplicate/empty rows out of {before}.)\")\n",
+    "    return df.reset_index(drop=True)\n",
+    "\n",
+    "        \n",
+    "# Search for posts containing a specific keyword\n",
+    "def get_posts_by_keyword(keyword, since, until, limit=300, page_size=100, max_retries=5):\n",
+    "\n",
+    "    state = _STATE.setdefault(keyword, {'until': until, 'rows': []})\n",
+    " \n",
+    "    if state['until'] != until and state['rows']:\n",
+    "        print(f\"Resuming '{keyword}' from until={state['until']}, {len(state['rows'])} posts already fetched this session.\")\n",
+    " \n",
+    "    since_dt = datetime.fromisoformat(since.replace('Z', '+00:00'))\n",
+    "    progress_bar = tqdm(total=limit, initial=min(len(state['rows']), limit), desc=\"Fetching posts\")\n",
+    "\n",
+    "    try:\n",
+    "        while len(state['rows']) < limit:\n",
+    "            try:\n",
+    "                search_results = call_with_retry(client.app.bsky.feed.search_posts, {'q': keyword, 'since': since, 'until': state['until'], 'lang': 'en', 'limit': page_size}, max_retries=max_retries)\n",
+    "            except Exception as e:\n",
+    "                print(f\"Giving up on this page after {max_retries} retries: {e}\")\n",
+    "                print(\"Progress so far is kept — just call this again to continue.\")\n",
+    "                break\n",
+    "    \n",
+    "            posts = search_results.posts\n",
+    "            if not posts:\n",
+    "                print(\"No more posts returned — reached the end of available results.\")\n",
+    "                break\n",
+    "    \n",
+    "            oldest_date_obj = None\n",
+    "            for post in posts:\n",
+    "                result = extract_post(post)\n",
+    "                if result is None:\n",
+    "                    continue\n",
+    "                row, date_obj = result\n",
+    "                state['rows'].append(row)\n",
+    "                if oldest_date_obj is None or date_obj < oldest_date_obj:\n",
+    "                    oldest_date_obj = date_obj\n",
+    "    \n",
+    "            progress_bar.update(len(posts))\n",
+    "    \n",
+    "            if oldest_date_obj is not None:\n",
+    "                state['until'] = oldest_date_obj.isoformat().replace('+00:00', 'Z')\n",
+    "                progress_bar.set_postfix_str(f'now at: {state[\"until\"]}')\n",
+    "    \n",
+    "                if oldest_date_obj < since_dt + timedelta(hours=1):\n",
+    "                    print(\"Reached the `since` boundary.\")\n",
+    "                    break\n",
+    " \n",
+    "    finally:\n",
+    "        progress_bar.close()\n",
+    " \n",
+    "    print(f\"{len(state['rows'])} posts fetched in total this session for '{keyword}'. Call again with the same checkpoint_name to keep going.\")\n",
+    "    return _dedupe(state['rows'])\n",
     "\n",
     "# Obtain posts from uri list\n",
-    "def get_posts(uris):\n",
+    "def get_posts(uris, chunk_size=20, max_retries=5):\n",
+    "    state = _STATE.setdefault('_by_uri', {'done_uris': set(), 'rows': []})\n",
+    " \n",
+    "    remaining = [u for u in uris if u not in state['done_uris']]\n",
+    "    chunks = list(chunked(remaining, chunk_size))\n",
+    " \n",
+    "    progress_bar = tqdm(total=len(uris), initial=len(state['done_uris']), desc=\"Fetching posts\")\n",
+    " \n",
     "    try:\n",
-    "        post_data = []\n",
+    "        for uribit in chunks:\n",
-    "\n",
+    "            try:\n",
-    "        progress_bar = tqdm(total=len(uris), desc=\"Fetching posts\")\n",
+    "                search_results = call_with_retry(client.app.bsky.feed.get_posts, {'uris': uribit}, max_retries=max_retries)\n",
-    "\n",
+    "            except Exception as e:\n",
-    "        for uribit in list(chunked(uris,20)):\n",
+    "                print(f\"Skipping this chunk of {len(uribit)} URIs after {max_retries} retries: {e}\")\n",
-    "\n",
-    "            search_results = client.app.bsky.feed.get_posts({'uris': uribit})\n",
-    "            posts = search_results.posts\n",
-    "\n",
-    "            for post in posts: \n",
-    "                # Extract post details\n",
-    "                text = clean_text(post.record.text) if hasattr(post.record, 'text') else np.nan\n",
-    "                iso_date = post.record.created_at if hasattr(post.record, 'created_at') else np.nan\n",
-    "                user = post.author.handle if hasattr(post.author, 'handle') else np.nan\n",
-    "\n",
-    "                post_id = post.uri if hasattr(post, 'uri') else np.nan\n",
-    "                parent_id = post.record.reply.parent.uri if hasattr(post.record, 'reply') and hasattr(post.record.reply, 'parent') else 'no_parent'\n",
-    "\n",
-    "                # Format date-time\n",
-    "                if iso_date != np.nan:\n",
-    "                    date_obj = datetime.fromisoformat(iso_date.replace('Z', '+00:00'))\n",
-    "                    date_str = date_obj.strftime('%Y-%m-%d %H:%M:%S')\n",
-    "                else:\n",
-    "                    date_str = np.nan\n",
-    "\n",
-    "                if post_id != np.nan and user != np.nan:\n",
-    "                    post_data.append({\n",
-    "                        'date': date_str,\n",
-    "                        'text': text,\n",
-    "                        'user': user,\n",
-    "                        'id': post_id,\n",
-    "                        'parentid': parent_id\n",
-    "                    })\n",
-    "                else:\n",
     "                continue\n",
     " \n",
-    "            progress_bar.update(len(posts))\n",
+    "            for post in search_results.posts:\n",
+    "                result = extract_post(post)\n",
+    "                if result is None:\n",
+    "                    continue\n",
+    "                row, _ = result\n",
+    "                state['rows'].append(row)\n",
     " \n",
+    "            state['done_uris'].update(uribit)\n",
+    "            progress_bar.update(len(uribit))\n",
+    " \n",
+    "    finally:\n",
     "        progress_bar.close()\n",
-    "        df = pd.DataFrame(post_data)\n",
+    " \n",
-    "        return df\n",
+    "    print(f\"{len(state['done_uris'])}/{len(uris)} URIs processed in total.\")\n",
-    "\n",
+    "    return _dedupe(state['rows'])"
-    "    except Exception as e:\n",
+   ]
-    "        print(f\"Error fetching posts: {e}\")\n",
+  },
-    "        return pd.DataFrame()"
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Define state"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "_STATE = {}"
    ]
   },
   {
@@ -233,7 +288,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "raw_df_data = get_posts_by_keyword('...', iso_start_timestamp, iso_end_timestamp, 500000)\n",
+    "raw_df_data = get_posts_by_keyword('...', iso_start_timestamp, iso_end_timestamp, 100000)\n",
     "raw_df_data.info()"
    ]
   },
@@ -251,13 +306,19 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "_STATE = {}\n",
     "id_list = raw_df_data['id'].to_list()\n",
     "id_list.append('no_parent')\n",
-    "\n",
+    "absent_parents = raw_df_data.loc[~raw_df_data['parentid'].isin(id_list), 'parentid'].tolist()"
-    "absent_parents = raw_df_data.loc[~raw_df_data['parentid'].isin(id_list), 'parentid'].tolist()\n",
+   ]
-    "\n",
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
     "df_absent_data = get_posts(absent_parents)\n",
-    "\n",
     "df_absent_data.info()"
    ]
   },
@@ -275,44 +336,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "df_data = pd.concat(raw_df_data, df_absent_data, ignore_index=True)\n",
+    "df_data = pd.concat([raw_df_data, df_absent_data], ignore_index=True)\n",
-    "\n",
-    "df_data.info()"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Drop duplicates."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "df_data = df_data.drop_duplicates()\n",
-    "df_data.info()"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Remove empty entries."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "df_data.dropna(inplace=True,subset=['text'])\n",
     "df_data.info()"
    ]
   },
@@ -356,7 +380,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "df_data.to_csv(f'../datasets/{case}.csv', index=False)"
+    "df_data.to_csv(f'{case}.csv', index=False)"
    ]
   },
   {
@@ -458,7 +482,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "closed_df_reduced_data.to_csv(f'../datasets/{case}-social-closed.csv', index=False)"
+    "closed_df_reduced_data.to_csv(f'{case}-social-closed.csv', index=False)"
    ]
   },
   {
@@ -475,7 +499,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "closed_df_data.to_csv(f'../datasets/{case}-closed.csv', index=False)"
+    "closed_df_data.to_csv(f'{case}-closed.csv', index=False)"
    ]
   }
  ],

graph-construct/construct-graph.ipynb

@@ -30,7 +30,7 @@
     "case = '...'\n",
     "casegraphs = '...'\n",
     "\n",
-    "df_data = pd.read_csv(f'../datasets/{case}.csv')\n",
+    "df_data = pd.read_csv(f'../sentiment-analysis/{case}.csv')\n",
     "user_nodes = df_data.groupby('user')['id'].apply(list).to_dict()\n",
     "\n",
     "df_data.info()"
@@ -62,7 +62,7 @@
     "        G.add_edge(row['id'], row['parentid'])\n",
     "\n",
     "# Export nearest-neighbour graph\n",
-    "nx.write_graphml(G, f'../datasets/{casegraphs}-nn-graph.graphml')\n",
+    "nx.write_graphml(G, f'{casegraphs}-nn-graph.graphml')\n",
     "\n",
     "# Define nearest-neighbour interaction count dictionary\n",
     "nn_count = {}\n",
@@ -98,7 +98,7 @@
     "tc_G = nx.transitive_closure(G)\n",
     "\n",
     "# Export transitive-closure graph\n",
-    "nx.write_graphml(tc_G, f'../datasets/{casegraphs}-tc-graph.graphml')\n",
+    "nx.write_graphml(tc_G, f'{casegraphs}-tc-graph.graphml')\n",
     "\n",
     "# Define transitive-closure interaction count dictionary\n",
     "tc_count = {}\n",
@@ -139,19 +139,19 @@
     "    if row['parentid'] != 'no_parent':\n",
     "        dG.add_edge(row['id'], row['parentid'])\n",
     "\n",
-    "nx.write_graphml(dG, f'../datasets/{casegraphs}-nn-dgraph.graphml')\n",
+    "nx.write_graphml(dG, f'{casegraphs}-nn-dgraph.graphml')\n",
     "\n",
     "# Construct directed-transitive-closure directed graph\n",
     "dtc_dG = nx.transitive_closure(dG)\n",
     "\n",
     "# Export directed-transitive-closure directed graph\n",
-    "nx.write_graphml(dtc_dG, f'../datasets/{casegraphs}-dtc-dgraph.graphml')\n",
+    "nx.write_graphml(dtc_dG, f'{casegraphs}-dtc-dgraph.graphml')\n",
     "\n",
     "# Convert to directed-transitive-closure undirected graph\n",
     "dtc_G = dtc_dG.to_undirected()\n",
     "\n",
     "# Export directed-transitive-closure graph\n",
-    "nx.write_graphml(dtc_G, f'../datasets/{casegraphs}-dtc-graph.graphml')\n",
+    "nx.write_graphml(dtc_G, f'{casegraphs}-dtc-graph.graphml')\n",
     "\n",
     "# Define directed-transitive-closure interaction count dictionary\n",
     "dtc_count = {}\n",
@@ -196,7 +196,7 @@
     "    'post-count': pd.Series(user_c)\n",
     "})\n",
     "\n",
-    "export_dataset = f'../datasets/{case}-users.csv'\n",
+    "export_dataset = f'{case}-users.csv'\n",
     "df_users.to_csv(export_dataset, index=False)\n",
     "\n",
     "df_users.info()"

graph-inspect/inspect-graph-comparison.ipynb

@@ -33,8 +33,8 @@
     "case2 = '...'\n",
     "\n",
     "# Import the graph\n",
-    "G1 = nx.read_graphml(f'../datasets/{case1}-dtc-graph.graphml')\n",
+    "G1 = nx.read_graphml(f'../graph-construct/{case1}-dtc-graph.graphml')\n",
-    "G2 = nx.read_graphml(f'../datasets/{case2}-dtc-graph.graphml')\n",
+    "G2 = nx.read_graphml(f'../graph-construct/{case2}-dtc-graph.graphml')\n",
     "\n",
     "# Print nodes, edges and density of the graphs\n",
     "print(f'G1:\\nNodes: {G1.number_of_nodes()}\\nEdges: {G1.number_of_edges()}\\nDensity: {nx.density(G1):.5f}')\n",

sentiment-analysis/sentiment-analysis.py

@@ -1,30 +1,56 @@
 import os
 import pandas as pd
-import torch
-from transformers import DistilBertTokenizerFast
-from torch.utils.data import DataLoader
-from transformers import DistilBertForSequenceClassification
 from tqdm import tqdm
+import torch
+from torch.utils.data import DataLoader, TensorDataset
+from transformers import (
+    DistilBertTokenizerFast,
+    DistilBertForSequenceClassification
+)
 
 # Importing the csv dataset.
 case = '...'
-df_data = pd.read_csv(f'../datasets/{case}.csv')
+df_data = pd.read_csv(f'{case}.csv')
 print(len(df_data))
 
-# Preparing the dataset for the network this includes tokenization, encoding and creating dataloaders.
+# Set device cuda
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+print('Using:', device)
+
+# Define model
 model_name = 'distilbert-base-uncased'
 tokenizer = DistilBertTokenizerFast.from_pretrained(model_name)
 
-encodings = tokenizer(df_data['text'].tolist(), truncation=True, padding=True, return_tensors='pt')
+# Tokenize and encode the text data
+encodings = tokenizer(
+    df_data['text'].tolist(), 
+    truncation=True, 
+    padding=True, 
+    return_tensors='pt'
+)
 
-dataset = torch.utils.data.TensorDataset(
+# Create tensor dataset
+dataset = TensorDataset(
     encodings['input_ids'], 
     encodings['attention_mask']
 )
-dataloader = DataLoader(dataset, batch_size=45, shuffle=False)
 
-# Loading the model.
+# Create data loader
-model= torch.load('models/bert-aclimdb.pth')
+if device == 'cuda':
+    en_pin_memory = True
+else:
+    en_pin_memory = False
+dataloader = DataLoader(
+    dataset, 
+    batch_size=53, 
+    shuffle=False,
+    pin_memory=en_pin_memory
+)
+
+# Loading the model
+model = DistilBertForSequenceClassification.from_pretrained(model_name, num_labels=1)
+model.load_state_dict(torch.load('...', map_location=device))
+model.to(device)
 
 # Using the model to perform sentiment analysis on the dataset.
 model.eval()
@@ -32,7 +58,7 @@ list_predicted_scores = []
 
 for batch in tqdm(dataloader):
     with torch.no_grad():
-        input_ids, attention_mask = batch
+        input_ids, attention_mask = [x.to(device) for x in batch]
 
         # Obtaining the sentiment score.
         output = model(input_ids=input_ids, attention_mask=attention_mask)
@@ -45,4 +71,4 @@ for batch in tqdm(dataloader):
 df_data['s'] = list_predicted_scores
 
 # Exporting the dataframe to a csv dataset.
-df_data.to_csv(f'../datasets/{case}-s.csv', index=False)
+df_data.to_csv(f'{case}-s.csv', index=False)

sentiment-analysis/training/aclimdb-sentiment-training.ipynb

@@ -27,8 +27,8 @@
     "import os\n",
     "import pandas as pd\n",
     "\n",
-    "imdb_train_dataset = \"aclimdb/train\"\n",
+    "imdb_train_dataset = 'aclimdb/train'\n",
-    "imdb_test_dataset = \"aclimdb/test\"\n",
+    "imdb_test_dataset = 'aclimdb/test'\n",
     "\n",
     "train_reviews = []\n",
     "train_scores = []\n",
@@ -118,7 +118,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "print(f\"Length training set: {len(df_train)}\\nLength testing set: {len(df_test)}\")"
+    "print(f'Length training set: {len(df_train)}\\nLength testing set: {len(df_test)}')"
    ]
   },
   {
@@ -136,31 +136,67 @@
    "outputs": [],
    "source": [
     "import torch\n",
-    "from transformers import DistilBertTokenizerFast\n",
+    "from torch.utils.data import DataLoader, TensorDataset\n",
-    "from torch.utils.data import DataLoader\n",
+    "from transformers import (\n",
-    "from transformers import DistilBertForSequenceClassification\n",
+    "    DistilBertTokenizerFast,\n",
+    "    DistilBertForSequenceClassification\n",
+    ")\n",
     "\n",
-    "model_name = \"distilbert-base-uncased\"\n",
+    "# Set device cuda\n",
+    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
+    "print('Using:', device)\n",
+    "\n",
+    "# Define model\n",
+    "model_name = 'distilbert-base-uncased'\n",
     "tokenizer = DistilBertTokenizerFast.from_pretrained(model_name)\n",
     "\n",
     "# Tokenize and encode the text data\n",
-    "train_encodings = tokenizer(df_train['text'].tolist(), truncation=True, padding=True, return_tensors='pt')\n",
+    "train_encodings = tokenizer(\n",
-    "test_encodings = tokenizer(df_test['text'].tolist(), truncation=True, padding=True, return_tensors='pt')\n",
+    "    df_train['text'].tolist(), \n",
+    "    truncation=True, \n",
+    "    padding=True, \n",
+    "    return_tensors='pt'\n",
+    ")\n",
     "\n",
-    "# Create data loaders\n",
+    "test_encodings = tokenizer(\n",
-    "train_dataset = torch.utils.data.TensorDataset(\n",
+    "    df_test['text'].tolist(), \n",
+    "    truncation=True, \n",
+    "    padding=True, \n",
+    "    return_tensors='pt'\n",
+    ")\n",
+    "\n",
+    "# Create tensor dataset\n",
+    "train_dataset = TensorDataset(\n",
     "    train_encodings['input_ids'], \n",
     "    train_encodings['attention_mask'], \n",
     "    torch.tensor(df_train['s'].values, dtype=torch.float32)\n",
     ")\n",
-    "train_dataloader = DataLoader(train_dataset, batch_size=25, shuffle=True)\n",
     "\n",
-    "test_dataset = torch.utils.data.TensorDataset(\n",
+    "test_dataset = TensorDataset(\n",
     "    test_encodings['input_ids'], \n",
     "    test_encodings['attention_mask'], \n",
     "    torch.tensor(df_test['s'].values, dtype=torch.float32)\n",
     ")\n",
-    "test_dataloader = DataLoader(test_dataset, batch_size=25, shuffle=False)"
+    "\n",
+    "# Create data loaders\n",
+    "if device == 'cuda':\n",
+    "    en_pin_memory = True\n",
+    "else:\n",
+    "    en_pin_memory = False\n",
+    "\n",
+    "train_dataloader = DataLoader(\n",
+    "    train_dataset, \n",
+    "    batch_size=10, \n",
+    "    shuffle=True,\n",
+    "    pin_memory=en_pin_memory\n",
+    ")\n",
+    "\n",
+    "test_dataloader = DataLoader(\n",
+    "    test_dataset, \n",
+    "    batch_size=10, \n",
+    "    shuffle=False,\n",
+    "    pin_memory=en_pin_memory\n",
+    ")"
    ]
   },
   {
@@ -177,7 +213,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "model = DistilBertForSequenceClassification.from_pretrained(model_name, num_labels=1)"
+    "model = DistilBertForSequenceClassification.from_pretrained(\n",
+    "    model_name, \n",
+    "    num_labels=1\n",
+    ").to(device)"
    ]
   },
   {
@@ -235,7 +274,7 @@
     "    for batch in train_dataloader:\n",
     "        global_step += 1\n",
     "        num_batches += 1\n",
-    "        input_ids, attention_mask, target_scores = batch\n",
+    "        input_ids, attention_mask, target_scores = [x.to(device) for x in batch]\n",
     "\n",
     "        # Forward pass\n",
     "        output = model(input_ids=input_ids, attention_mask=attention_mask)\n",
@@ -270,7 +309,7 @@
     "\n",
     "    for batch in test_dataloader:\n",
     "        with torch.no_grad():\n",
-    "            input_ids, attention_mask, target_scores = batch\n",
+    "            input_ids, attention_mask, target_scores = [x.to(device) for x in batch]\n",
     "\n",
     "            # Obtaining the scores\n",
     "            output = model(input_ids=input_ids, attention_mask=attention_mask)\n",
@@ -290,7 +329,7 @@
     "    print(f\"Epoch {epoch + 1}/{num_epochs}, Validation loss: {average_loss:.4f}\\n\")\n",
     "\n",
     "    # Saving the model\n",
-    "    torch.save(model, f'bert-aclimdb/{epoch + 1}.pth')\n",
+    "    torch.save(model.state_dict(), f'bert-aclimdb/{epoch + 1}.pth')\n",
     "\n",
     "    # Early stopping check\n",
     "    if average_loss < best_validation_loss:\n",
@@ -300,6 +339,7 @@
     "        no_improvement_counter += 1\n",
     "\n",
     "    if no_improvement_counter >= early_stop_patience:\n",
+    "        print('Early stopping triggered')\n",
     "        break\n",
     "\n",
     "writer.close()"
@@ -319,7 +359,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "model = torch.load('bert-aclimdb/4.pth')"
+    "model = DistilBertForSequenceClassification.from_pretrained(model_name, num_labels=1)\n",
+    "model.load_state_dict(torch.load('...', map_location=device))\n",
+    "model.to(device)"
    ]
   },
   {
@@ -344,7 +386,7 @@
     "\n",
     "for batch in test_dataloader:\n",
     "    with torch.no_grad():\n",
-    "        input_ids, attention_mask, scores = batch\n",
+    "        input_ids, attention_mask, scores = [x.to(device) for x in batch]\n",
     "\n",
     "        # Obtaining the scores\n",
     "        output = model(input_ids=input_ids, attention_mask=attention_mask)\n",