ai-nk-cce 0.1.0__py3-none-any.whl

model/score.ipynb

@@ -0,0 +1,596 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from datasets import load_from_disk\n",
+    "from src.model.parser import eval_target, target_to_int\n",
+    "import matplotlib.pyplot as plt\n",
+    "import seaborn as sns\n",
+    "\n",
+    "# dataset_file = \"../data/model_evals/gpt2_s10_v2_lr_scan_1/1e-5\"\n",
+    "# dataset_file = \"../data/model_evals/gpt2_s10_lr_scan_v3/1e-5\"\n",
+    "dataset_file = \"/u/lumi/projects/human-ai-social-learning/models/gpt2_v6/rl/2025_07_15__13_51_57/inference_results_xxl\"\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ds = load_from_disk(dataset_file)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Retrieve rank of suggestion\n",
+    "\n",
+    "To retrieve the rank of the suggestion, we extract the rank list and look which value the int version of suggestion has in it."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "\n",
+    "def eval_row(row):\n",
+    "    try:\n",
+    "        suggestion_rank = eval_target(row['suggestion'], row['ranks'])\n",
+    "    except Exception as e:\n",
+    "        print(f\"Error evaluating row: {e}\")\n",
+    "        suggestion_rank = -1\n",
+    "    return {\n",
+    "        'suggestion_rank': suggestion_rank,\n",
+    "    }"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ds['test'] = ds['test'].map(eval_row)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "*Extra analysis*:\n",
+    "\n",
+    "Below we found out, that a lot of suggestions are having rank -1. Let's find out, if those are because of suggesting the user input."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "df_neg_one = ds['test'].to_pandas().query('suggestion_rank == -1')\n",
+    "df_neg_one['suggested_origin'] =  (df_neg_one['suggestion'].apply(target_to_int) == df_neg_one['origin_idx'])\n",
+    "df_neg_one['n_samples'] = df_neg_one['context'].str.count(r'[01](?:,[01]){7}')-1\n",
+    "\n",
+    "# Calculate percentage of suggested_origin for each combination of n_samples and hamming_distance\n",
+    "suggested_origin_pivot = df_neg_one.pivot_table(\n",
+    "    values='suggested_origin',\n",
+    "    index='hamming_distance',\n",
+    "    columns='n_samples',\n",
+    "    aggfunc=lambda x: (x.sum() / len(x)) * 100\n",
+    ")\n",
+    "print(\"\\nPercentage of -1 ranks that are suggested origin:\")\n",
+    "print(suggested_origin_pivot)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Get necessary information to plot\n",
+    "\n",
+    "Extract rank, hamming_dist, k and the amount of samples from the ds into a df."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "df = ds['test'].to_pandas()[['suggestion', 'suggestion_rank', 'hamming_distance', 'k', 'context', 'ranks', 'origin_idx', 'payoffs']]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Count the amount of samples, by regexing."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "df['n_samples'] = df['context'].str.count(r'[01](?:,[01]){7}')-1"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Extracting average suggestion rank\n",
+    "\n",
+    "We group and aggregate the average suggestion rank for *hamming_distance*, *k* and *n_samples*."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Group by n_constraints and k, then calculate mean suggestion_rank\n",
+    "agg_df = df.groupby(['hamming_distance', 'k', 'n_samples'])['suggestion_rank'].mean().reset_index()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "*Extra analysis*: As we saw a lot of *suggestion_rank* == -1. We want to know how many."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Calculate percentage of -1 ranks for each combination of hamming_distance and n_samples\n",
+    "neg_one_df = df.pivot_table(\n",
+    "    values='suggestion_rank',\n",
+    "    index='hamming_distance', \n",
+    "    columns='n_samples',\n",
+    "    aggfunc=lambda x: (x == -1).mean() * 100\n",
+    ").round(2)\n",
+    "\n",
+    "print(\"Percentage of suggestion_rank == -1 by hamming_distance and n_samples:\")\n",
+    "print(neg_one_df)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Plotting results\n",
+    "\n",
+    "The plot depicts the average of *suggestion_rank* for a combination of either *hamming_distance*, *n_samples* and/or *k*."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Create a pivot table with hamming_distance as rows and k as columns\n",
+    "agg_df = df.pivot_table(\n",
+    "    values='suggestion_rank',\n",
+    "    index='hamming_distance',\n",
+    "    columns='k',\n",
+    "    aggfunc='mean'\n",
+    ")\n",
+    "\n",
+    "# Round the values to 3 decimal places for better readability\n",
+    "agg_df = agg_df.round(3)\n",
+    "\n",
+    "# Create a figure with a larger size for better readability\n",
+    "plt.figure(figsize=(12, 8))\n",
+    "\n",
+    "# Create the heatmap with diverging colormap centered at 0.5\n",
+    "sns.heatmap(\n",
+    "    agg_df,\n",
+    "    annot=True,\n",
+    "    cmap='RdBu_r',  # Red-blue diverging colormap (reversed)\n",
+    "    center=0.5,     # Center the colormap at 0.5 for clear above/below contrast\n",
+    "    fmt='.3f',\n",
+    "    cbar_kws={'label': 'Average Suggestion Rank'},\n",
+    "    square=True\n",
+    ")\n",
+    "\n",
+    "# Customize the plot\n",
+    "plt.title('Average Suggestion Rank by Hamming Distance and K', pad=20)\n",
+    "plt.xlabel('K Value')\n",
+    "plt.ylabel('Hamming Distance')\n",
+    "\n",
+    "# Adjust layout to prevent label cutoff\n",
+    "plt.tight_layout()\n",
+    "\n",
+    "# Show the plot\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Calculate improvement_mean for each row\n",
+    "def calculate_improvement_mean(row):\n",
+    "    try:\n",
+    "        # Get user's original choice payoff\n",
+    "        origin_payoff = row['payoffs'][row['origin_idx']]\n",
+    "        \n",
+    "        # Convert suggestion to integer index and get its payoff\n",
+    "        suggestion_idx = int(row['suggestion'].replace(',', ''), 2)\n",
+    "        suggestion_payoff = row['payoffs'][suggestion_idx]\n",
+    "        \n",
+    "        # Calculate improvement as difference (suggestion - origin)\n",
+    "        # Positive means suggestion is better, negative means origin was better\n",
+    "        improvement_mean = (suggestion_payoff - origin_payoff) / 1000\n",
+    "        return improvement_mean\n",
+    "    except:\n",
+    "        return None\n",
+    "\n",
+    "# Add improvement_mean to the dataframe\n",
+    "df['improvement_mean'] = df.apply(calculate_improvement_mean, axis=1)\n",
+    "\n",
+    "print(f\"\\nImprovement_mean statistics:\")\n",
+    "print(f\"Min: {df['improvement_mean'].min():.3f}\")\n",
+    "print(f\"Max: {df['improvement_mean'].max():.3f}\")\n",
+    "print(f\"Mean: {df['improvement_mean'].mean():.3f}\")\n",
+    "print(f\"Median: {df['improvement_mean'].median():.3f}\")\n",
+    "\n",
+    "# Check if we have values above and below some meaningful threshold\n",
+    "print(f\"\\nValues below 0: {(df['improvement_mean'] < 0).sum()}\")\n",
+    "print(f\"Values above 0: {(df['improvement_mean'] > 0).sum()}\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Create a pivot table to count the number of values in each cell\n",
+    "count_df = df.pivot_table(\n",
+    "    values='improvement_mean',\n",
+    "    index='hamming_distance', \n",
+    "    columns='k',\n",
+    "    aggfunc='count'\n",
+    ")\n",
+    "\n",
+    "print(\"Number of values in each cell of the pivot table:\")\n",
+    "print(count_df)\n",
+    "print(\"\\n\")\n",
+    "\n",
+    "# Create a pivot table with hamming_distance as rows and k as columns\n",
+    "# using improvement_mean instead of suggestion_rank\n",
+    "agg_df = df.pivot_table(\n",
+    "    values='improvement_mean',\n",
+    "    index='hamming_distance',\n",
+    "    columns='k',\n",
+    "    aggfunc='mean'\n",
+    ")\n",
+    "\n",
+    "# Round the values to 3 decimal places for better readability\n",
+    "agg_df = agg_df.round(3)\n",
+    "\n",
+    "# Create a figure with a larger size for better readability\n",
+    "plt.figure(figsize=(12, 8))\n",
+    "\n",
+    "# Create the heatmap with diverging colormap centered at 0\n",
+    "sns.heatmap(\n",
+    "    agg_df,\n",
+    "    annot=True,\n",
+    "    cmap='RdYlBu_r',  # Red-yellow-blue diverging colormap (reversed)\n",
+    "    center=0.0,       # Center the colormap at 0 for clear above/below contrast\n",
+    "    fmt='.3f',\n",
+    "    cbar_kws={'label': 'Average Improvement Mean'},\n",
+    "    square=True\n",
+    ")\n",
+    "\n",
+    "# Customize the plot\n",
+    "plt.title('Average Improvement Mean by Hamming Distance and K', \n",
+    "          pad=20)\n",
+    "plt.xlabel('K Value')\n",
+    "plt.ylabel('Hamming Distance')\n",
+    "\n",
+    "# Adjust layout to prevent label cutoff\n",
+    "plt.tight_layout()\n",
+    "\n",
+    "# Show the plot\n",
+    "plt.show()\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "vscode": {
+     "languageId": "raw"
+    }
+   },
+   "source": [
+    "## Distribution Analysis of Suggestion Ranks\n",
+    "\n",
+    "Analyze the distribution of suggestion_rank values grouped by hamming_distance and k.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pandas as pd\n",
+    "from collections import defaultdict\n",
+    "\n",
+    "# Group by hamming_distance and k, then collect suggestion_rank distributions\n",
+    "distributions = {}\n",
+    "\n",
+    "# Filter out -1 values for cleaner distribution analysis\n",
+    "df_clean = df[df['suggestion_rank'] != -1].copy()\n",
+    "\n",
+    "for (hamming_dist, k), group in df_clean.groupby(['hamming_distance', 'k']):\n",
+    "    suggestion_ranks = group['suggestion_rank'].values\n",
+    "    \n",
+    "    # Calculate distribution metrics\n",
+    "    dist_info = {\n",
+    "        'count': len(suggestion_ranks),\n",
+    "        'mean': suggestion_ranks.mean(),\n",
+    "        'std': suggestion_ranks.std(),\n",
+    "        'min': suggestion_ranks.min(),\n",
+    "        'max': suggestion_ranks.max(),\n",
+    "        'q25': np.percentile(suggestion_ranks, 25),\n",
+    "        'q50': np.percentile(suggestion_ranks, 50),  # median\n",
+    "        'q75': np.percentile(suggestion_ranks, 75),\n",
+    "        'values': suggestion_ranks  # Keep all values for plotting\n",
+    "    }\n",
+    "    \n",
+    "    distributions[(hamming_dist, k)] = dist_info\n",
+    "\n",
+    "# Display summary of distributions\n",
+    "print(\"Distribution summary (hamming_distance, k): count, mean ± std\")\n",
+    "for (h_dist, k), info in sorted(distributions.items()):\n",
+    "    print(f\"({h_dist}, {k}): n={info['count']:2d}, \"\n",
+    "          f\"μ={info['mean']:.2f}±{info['std']:.2f}, \"\n",
+    "          f\"range=[{info['min']:.1f}, {info['max']:.1f}]\")\n",
+    "\n",
+    "print(f\"\\nTotal combinations: {len(distributions)}\")\n",
+    "print(f\"Sample sizes range: {min(d['count'] for d in distributions.values())} \"\n",
+    "      f\"to {max(d['count'] for d in distributions.values())}\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Set up beautiful plot styling\n",
+    "plt.style.use('seaborn-v0_8-whitegrid')\n",
+    "sns.set_palette(\"husl\")\n",
+    "\n",
+    "# Define beautiful color palette\n",
+    "COLORS = {\n",
+    "    'primary': '#2E86AB',      # Deep blue\n",
+    "    'secondary': '#A23B72',    # Deep pink\n",
+    "    'accent': '#F18F01',       # Orange\n",
+    "    'success': '#C73E1D',      # Red\n",
+    "    'histogram': '#4A90E2',    # Light blue\n",
+    "    'boxplot': '#4A90E2',      # Deep blue\n",
+    "    'scatter_main': '#2E86AB', # Deep blue\n",
+    "    'scatter_points': '#87CEEB', # Sky blue\n",
+    "    'error_bars': '#FF6B6B',   # Coral red\n",
+    "    'grid': '#E8E8E8',         # Light gray\n",
+    "    'text': '#2C3E50'          # Dark blue-gray\n",
+    "}\n",
+    "\n",
+    "def plot_distribution(hamming_distance, k, plot_type='all'):\n",
+    "    \"\"\"\n",
+    "    Plot the distribution of suggestion_rank showing histogram, boxplot, \n",
+    "    and scatter plot with error bars side-by-side with beautiful styling.\n",
+    "    \n",
+    "    Parameters:\n",
+    "    -----------\n",
+    "    hamming_distance : int\n",
+    "        The hamming distance value\n",
+    "    k : int\n",
+    "        The k value\n",
+    "    plot_type : str\n",
+    "        Type of plot: 'all' (histogram, boxplot, and scatter), \n",
+    "        'histogram', 'boxplot', or 'scatter'\n",
+    "    \"\"\"\n",
+    "    key = (hamming_distance, k)\n",
+    "    \n",
+    "    if key not in distributions:\n",
+    "        print(f\"No data found for hamming_distance={hamming_distance}, \"\n",
+    "              f\"k={k}\")\n",
+    "        available_keys = sorted(list(distributions.keys()))\n",
+    "        print(\"Available combinations:\")\n",
+    "        for h_dist, k_val in available_keys:\n",
+    "            print(f\"  hamming_distance={h_dist}, k={k_val}\")\n",
+    "        return\n",
+    "    \n",
+    "    dist_info = distributions[key]\n",
+    "    values = dist_info['values']\n",
+    "    \n",
+    "    # Create figure based on plot_type\n",
+    "    if plot_type == 'all':\n",
+    "        fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(20, 6))\n",
+    "        axes = [ax1, ax2, ax3]\n",
+    "        plot_types = ['histogram', 'boxplot', 'scatter']\n",
+    "    elif plot_type in ['histogram', 'boxplot', 'scatter']:\n",
+    "        fig, ax = plt.subplots(1, 1, figsize=(8, 6))\n",
+    "        axes = [ax]\n",
+    "        plot_types = [plot_type]\n",
+    "    else:\n",
+    "        print(f\"Invalid plot_type: {plot_type}\")\n",
+    "        print(\"Available types: 'all', 'histogram', 'boxplot', 'scatter'\")\n",
+    "        return\n",
+    "    \n",
+    "    # Set overall figure background and styling\n",
+    "    fig.patch.set_facecolor('white')\n",
+    "    \n",
+    "    # Plot each type with beautiful styling\n",
+    "    for ax, ptype in zip(axes, plot_types):\n",
+    "        if ptype == 'histogram':\n",
+    "            # Beautiful histogram\n",
+    "            bins = min(15, len(np.unique(values)))\n",
+    "            n, bins_edges, patches = ax.hist(values, \n",
+    "                                            bins=bins,\n",
+    "                                            alpha=0.8,\n",
+    "                                            color=COLORS['histogram'],\n",
+    "                                            edgecolor='white',\n",
+    "                                            linewidth=1.2,\n",
+    "                                            orientation='horizontal'\n",
+    "                                            )\n",
+    "            \n",
+    "            # Gradient effect for histogram bars\n",
+    "            for i, patch in enumerate(patches):\n",
+    "                patch.set_facecolor(plt.cm.Blues(0.4 + 0.4 * i / len(patches)))\n",
+    "            \n",
+    "            ax.set_xlabel('Count', fontsize=12, color=COLORS['text'])\n",
+    "            ax.set_ylabel('Suggestion Rank', fontsize=12, color=COLORS['text'])\n",
+    "            ax.set_title(f'Distribution Histogram\\nHamming Distance = {hamming_distance}, k = {k}',\n",
+    "                        fontsize=14, fontweight='bold', color=COLORS['text'], pad=20)\n",
+    "                            \n",
+    "        elif ptype == 'boxplot':\n",
+    "            # Beautiful boxplot\n",
+    "            bp = ax.boxplot(values, vert=True, patch_artist=True,\n",
+    "                           boxprops=dict(facecolor=COLORS['boxplot'], alpha=0.8,\n",
+    "                                       edgecolor=COLORS['text'], linewidth=1.5),\n",
+    "                           medianprops=dict(color='white', linewidth=2.5),\n",
+    "                           whiskerprops=dict(color=COLORS['text'], linewidth=1.5),\n",
+    "                           capprops=dict(color=COLORS['text'], linewidth=1.5),\n",
+    "                           flierprops=dict(marker='o', markerfacecolor=COLORS['accent'], \n",
+    "                                         markeredgecolor=COLORS['text'], markersize=6, alpha=0.7))\n",
+    "            \n",
+    "            ax.set_ylabel('Suggestion Rank', fontsize=12, color=COLORS['text'])\n",
+    "            ax.set_title(f'Distribution Boxplot\\nHamming Distance = {hamming_distance}, k = {k}', \n",
+    "                        fontsize=14, fontweight='bold', color=COLORS['text'], pad=20)\n",
+    "            ax.set_xticklabels([f'H={hamming_distance}\\nk={k}'], fontsize=11)\n",
+    "            \n",
+    "        elif ptype == 'scatter':\n",
+    "            # Beautiful scatter plot with error bars\n",
+    "            x_pos = hamming_distance\n",
+    "            y_mean = dist_info['mean']\n",
+    "            y_std = dist_info['std']\n",
+    "            \n",
+    "            # Plot mean with beautiful error bars\n",
+    "            ax.errorbar(x_pos, y_mean, yerr=y_std, \n",
+    "                       fmt='o', markersize=12, capsize=8, capthick=3,\n",
+    "                       color=COLORS['scatter_main'], ecolor=COLORS['error_bars'], \n",
+    "                       alpha=0.9, linewidth=2.5, markeredgecolor='white',\n",
+    "                       markeredgewidth=2)\n",
+    "            \n",
+    "            # Add individual data points with jitter and beautiful styling\n",
+    "            jitter = np.random.normal(0, 0.04, len(values))\n",
+    "            ax.scatter(x_pos + jitter, values, alpha=0.4, s=25, \n",
+    "                      color=COLORS['scatter_points'], \n",
+    "                      edgecolors=COLORS['scatter_main'], \n",
+    "                      linewidth=0.8)\n",
+    "            \n",
+    "            ax.set_xlabel('Hamming Distance', fontsize=12, color=COLORS['text'])\n",
+    "            ax.set_ylabel('Suggestion Rank', fontsize=12, color=COLORS['text'])\n",
+    "            ax.set_title(f'Scatter Plot with Error Bars\\nHamming Distance = {hamming_distance}, k = {k}', \n",
+    "                        fontsize=14, fontweight='bold', color=COLORS['text'], pad=20)\n",
+    "            \n",
+    "            # Set x-axis styling\n",
+    "            ax.set_xlim(hamming_distance - 0.6, hamming_distance + 0.6)\n",
+    "            ax.set_xticks([hamming_distance])\n",
+    "        \n",
+    "        # Common styling for all plots\n",
+    "        ax.grid(True, alpha=0.3, color=COLORS['grid'], linestyle='-', linewidth=0.8)\n",
+    "        ax.set_facecolor('#FAFAFA')\n",
+    "        ax.tick_params(colors=COLORS['text'], labelsize=10)\n",
+    "\n",
+    "\n",
+    "        # Beautiful statistics box\n",
+    "        stats_text = (f'n = {dist_info[\"count\"]:,}\\n'\n",
+    "                        f'μ = {dist_info[\"mean\"]:.3f}\\n'\n",
+    "                        f'σ = {dist_info[\"std\"]:.3f}\\n'\n",
+    "                        f'median = {dist_info[\"q50\"]:.3f}')\n",
+    "        ax.text(0.05, 0.95, stats_text, transform=ax.transAxes,\n",
+    "                verticalalignment='top', fontsize=11,\n",
+    "                bbox=dict(boxstyle=\"round,pad=0.5\", \n",
+    "                        facecolor='white', alpha=0.9,\n",
+    "                        edgecolor=COLORS['primary'], linewidth=2))\n",
+    "        \n",
+    "        # Beautiful spines\n",
+    "        for spine in ax.spines.values():\n",
+    "            spine.set_color(COLORS['text'])\n",
+    "            spine.set_linewidth(1.2)\n",
+    "    \n",
+    "    plt.tight_layout(pad=3.0)\n",
+    "    plt.show()\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Beautiful Plot Examples with New Styling\n",
+    "\n",
+    "print(\"=== Beautiful Distribution Plots ===\")\n",
+    "print(\"Available plot types: 'all', 'histogram', 'boxplot', 'scatter'\")\n",
+    "print()\n",
+    "\n",
+    "# Example 1: Show all three plots side-by-side (default)\n",
+    "print(\"1. All plots side-by-side (histogram, boxplot, scatter):\")\n",
+    "plot_distribution(hamming_distance=4, k=4, plot_type='all')\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "ai-nk-cce-ros7_1vM-py3.10",
+   "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.10.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

model/scripts/template.slurm

@@ -0,0 +1,54 @@
+#!/bin/bash -l
+#SBATCH --output {output_dir}/slurm-%x-%j.out
+#SBATCH --error {output_dir}/slurm-%x-%j.out
+#SBATCH --chdir ./
+#SBATCH --job-name {job_name}/{job_id}
+#
+#SBATCH --nodes={n_nodes}
+#SBATCH --tasks-per-node=1
+#SBATCH --cpus-per-task={n_cpu}
+#SBATCH --mem={memory}
+#
+#SBATCH --constraint="gpu"
+#SBATCH --gres=gpu:a100:{n_gpu}
+#SBATCH --partition=gpu
+
+# Wall clock limit (max is 24 hours):
+#SBATCH --time={time}
+
+module purge
+module load apptainer
+
+source .env
+
+# create huggingface cache directory if it doesn't exist
+mkdir -p ~/.cache/huggingface
+
+echo "Runing training using the image: {image}"
+echo "Runing training using the config: {config_file}"
+
+srun apptainer exec \
+	--nv \
+    --contain \
+    --cleanenv \
+    --pwd /root/llm-strategic-tuning \
+    --bind .:/root/llm-strategic-tuning \
+    --bind ~/.cache/huggingface:/root/.cache/huggingface \
+    --bind /ptmp:/ptmp \
+    --env HUGGING_FACE_HUB_TOKEN="$HUGGINGFACE_TOKEN" \
+    --env WANDB_API_KEY="$WANDB_API_KEY" \
+    --env WANDB_ENTITY="chm-ml" \
+    --env WANDB_PROJECT="{project_name}" \
+    --env WANDB_RUN_GROUP="{group_name}" \
+    --env WANDB_NAME="{job_name}/{job_id}" \
+    --env NCCL_DEBUG="INFO" \
+    --env NCCL_BLOCKING_WAIT="0" \
+    --env HF_HOME="/root/.cache/huggingface" \
+	{image} \
+    python -m torch.distributed.run \
+		--nnodes="$SLURM_NNODES" \
+		--nproc-per-node=gpu \
+		--rdzv-id="$SLURM_JOBID" \
+		--rdzv-endpoint=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1) \
+		--rdzv-backend="c10d" \
+    {script} --config {config_file}