Merge pull request #1 from jenbenarye/train

ml/dataset_training.ipynb

@@ -1,398 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": 43,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "#dependencies:\n",
-    "import pandas as pd\n",
-    "\n",
-    "import torch\n",
-    "from transformers import GPT2Tokenizer\n",
-    "\n",
-    "from trl import AutoModelForCausalLMWithValueHead, PPOConfig, PPOTrainer"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 44,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b8a22b8d60c0417eafbf554832398287",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Resolving data files:   0%|          | 0/18 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b83d2624c2b14986a8297821460225ab",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Resolving data files:   0%|          | 0/18 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b4304c0f48cb472589b5e80d3a42cba2",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Resolving data files:   0%|          | 0/18 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
-   "source": [
-    "#loading datasets:\n",
-    "from datasets import load_dataset\n",
-    "\n",
-    "ds = load_dataset(\"stanfordnlp/SHP\", split='train')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 45,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Index(['post_id', 'domain', 'upvote_ratio', 'history', 'c_root_id_A',\n",
-      "       'c_root_id_B', 'created_at_utc_A', 'created_at_utc_B', 'score_A',\n",
-      "       'score_B', 'human_ref_A', 'human_ref_B', 'labels', 'seconds_difference',\n",
-      "       'score_ratio'],\n",
-      "      dtype='object')\n"
-     ]
-    }
-   ],
-   "source": [
-    "df = ds.to_pandas()\n",
-    "print(df.columns)\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 46,
-   "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>upvote_ratio</th>\n",
-       "      <th>history</th>\n",
-       "      <th>score_A</th>\n",
-       "      <th>score_B</th>\n",
-       "      <th>human_ref_A</th>\n",
-       "      <th>human_ref_B</th>\n",
-       "      <th>labels</th>\n",
-       "      <th>score_ratio</th>\n",
-       "    </tr>\n",
-       "  </thead>\n",
-       "  <tbody>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>0.99</td>\n",
-       "      <td>In an interview right before receiving the 201...</td>\n",
-       "      <td>52</td>\n",
-       "      <td>54</td>\n",
-       "      <td>Currently wrapping up my PhD. There is a stark...</td>\n",
-       "      <td>It’s ironic to me that research has shown that...</td>\n",
-       "      <td>0</td>\n",
-       "      <td>1.038462</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1</th>\n",
-       "      <td>0.95</td>\n",
-       "      <td>If any professor is reading this: please do no...</td>\n",
-       "      <td>5</td>\n",
-       "      <td>17</td>\n",
-       "      <td>And when your teacher doesn't listen or pay at...</td>\n",
-       "      <td>I'm pretty strict on time, to the point where ...</td>\n",
-       "      <td>0</td>\n",
-       "      <td>3.400000</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>2</th>\n",
-       "      <td>0.95</td>\n",
-       "      <td>If any professor is reading this: please do no...</td>\n",
-       "      <td>5</td>\n",
-       "      <td>7</td>\n",
-       "      <td>Profs can be oblivious? What’s new!</td>\n",
-       "      <td>This sounds like a problem with a specific pro...</td>\n",
-       "      <td>0</td>\n",
-       "      <td>1.400000</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>3</th>\n",
-       "      <td>0.95</td>\n",
-       "      <td>If any professor is reading this: please do no...</td>\n",
-       "      <td>7</td>\n",
-       "      <td>5</td>\n",
-       "      <td>This sounds like a problem with a specific pro...</td>\n",
-       "      <td>And when your teacher doesn't listen or pay at...</td>\n",
-       "      <td>1</td>\n",
-       "      <td>1.400000</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>4</th>\n",
-       "      <td>0.95</td>\n",
-       "      <td>If any professor is reading this: please do no...</td>\n",
-       "      <td>6</td>\n",
-       "      <td>7</td>\n",
-       "      <td>This would be totally unacceptable in my class...</td>\n",
-       "      <td>This sounds like a problem with a specific pro...</td>\n",
-       "      <td>0</td>\n",
-       "      <td>1.166667</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>...</th>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>348713</th>\n",
-       "      <td>0.94</td>\n",
-       "      <td>Can I get in trouble for giving my neighbor hi...</td>\n",
-       "      <td>7</td>\n",
-       "      <td>25</td>\n",
-       "      <td>Just put up a fence. Legally he isn't responsi...</td>\n",
-       "      <td>Whatever you do, don't cut his trees down.</td>\n",
-       "      <td>0</td>\n",
-       "      <td>3.571429</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>348714</th>\n",
-       "      <td>0.94</td>\n",
-       "      <td>Can I get in trouble for giving my neighbor hi...</td>\n",
-       "      <td>2</td>\n",
-       "      <td>25</td>\n",
-       "      <td>If OP pays someone to clean his yard, and then...</td>\n",
-       "      <td>Whatever you do, don't cut his trees down.</td>\n",
-       "      <td>0</td>\n",
-       "      <td>12.500000</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>348715</th>\n",
-       "      <td>0.94</td>\n",
-       "      <td>Can I get in trouble for giving my neighbor hi...</td>\n",
-       "      <td>9</td>\n",
-       "      <td>7</td>\n",
-       "      <td>My  observation is that both of you are idiots...</td>\n",
-       "      <td>Are you Rand Paul's neighbor?  https://www.gq....</td>\n",
-       "      <td>1</td>\n",
-       "      <td>1.285714</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>348716</th>\n",
-       "      <td>0.94</td>\n",
-       "      <td>Can I get in trouble for giving my neighbor hi...</td>\n",
-       "      <td>9</td>\n",
-       "      <td>7</td>\n",
-       "      <td>My  observation is that both of you are idiots...</td>\n",
-       "      <td>Just put up a fence. Legally he isn't responsi...</td>\n",
-       "      <td>1</td>\n",
-       "      <td>1.285714</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>348717</th>\n",
-       "      <td>0.94</td>\n",
-       "      <td>Can I get in trouble for giving my neighbor hi...</td>\n",
-       "      <td>7</td>\n",
-       "      <td>2</td>\n",
-       "      <td>Capture his acts on camera.  Collect and bag l...</td>\n",
-       "      <td>If OP pays someone to clean his yard, and then...</td>\n",
-       "      <td>1</td>\n",
-       "      <td>3.500000</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table>\n",
-       "<p>348718 rows × 8 columns</p>\n",
-       "</div>"
-      ],
-      "text/plain": [
-       "        upvote_ratio                                            history  \\\n",
-       "0               0.99  In an interview right before receiving the 201...   \n",
-       "1               0.95  If any professor is reading this: please do no...   \n",
-       "2               0.95  If any professor is reading this: please do no...   \n",
-       "3               0.95  If any professor is reading this: please do no...   \n",
-       "4               0.95  If any professor is reading this: please do no...   \n",
-       "...              ...                                                ...   \n",
-       "348713          0.94  Can I get in trouble for giving my neighbor hi...   \n",
-       "348714          0.94  Can I get in trouble for giving my neighbor hi...   \n",
-       "348715          0.94  Can I get in trouble for giving my neighbor hi...   \n",
-       "348716          0.94  Can I get in trouble for giving my neighbor hi...   \n",
-       "348717          0.94  Can I get in trouble for giving my neighbor hi...   \n",
-       "\n",
-       "        score_A  score_B                                        human_ref_A  \\\n",
-       "0            52       54  Currently wrapping up my PhD. There is a stark...   \n",
-       "1             5       17  And when your teacher doesn't listen or pay at...   \n",
-       "2             5        7                Profs can be oblivious? What’s new!   \n",
-       "3             7        5  This sounds like a problem with a specific pro...   \n",
-       "4             6        7  This would be totally unacceptable in my class...   \n",
-       "...         ...      ...                                                ...   \n",
-       "348713        7       25  Just put up a fence. Legally he isn't responsi...   \n",
-       "348714        2       25  If OP pays someone to clean his yard, and then...   \n",
-       "348715        9        7  My  observation is that both of you are idiots...   \n",
-       "348716        9        7  My  observation is that both of you are idiots...   \n",
-       "348717        7        2  Capture his acts on camera.  Collect and bag l...   \n",
-       "\n",
-       "                                              human_ref_B  labels  score_ratio  \n",
-       "0       It’s ironic to me that research has shown that...       0     1.038462  \n",
-       "1       I'm pretty strict on time, to the point where ...       0     3.400000  \n",
-       "2       This sounds like a problem with a specific pro...       0     1.400000  \n",
-       "3       And when your teacher doesn't listen or pay at...       1     1.400000  \n",
-       "4       This sounds like a problem with a specific pro...       0     1.166667  \n",
-       "...                                                   ...     ...          ...  \n",
-       "348713         Whatever you do, don't cut his trees down.       0     3.571429  \n",
-       "348714         Whatever you do, don't cut his trees down.       0    12.500000  \n",
-       "348715  Are you Rand Paul's neighbor?  https://www.gq....       1     1.285714  \n",
-       "348716  Just put up a fence. Legally he isn't responsi...       1     1.285714  \n",
-       "348717  If OP pays someone to clean his yard, and then...       1     3.500000  \n",
-       "\n",
-       "[348718 rows x 8 columns]"
-      ]
-     },
-     "execution_count": 46,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "# df['response_length'] = df['history'].apply(len)\n",
-    "# df['label'] = df['response_length'].apply(lambda x: 'long' if x > 100 else 'short')\n",
-    "df.drop(columns=['post_id', 'domain', 'c_root_id_A', 'c_root_id_B', 'created_at_utc_A', 'created_at_utc_B', 'seconds_difference'])"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 47,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/Users/riddhib/.pyenv/versions/3.10.13/lib/python3.10/site-packages/transformers/tokenization_utils_base.py:1617: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be deprecated in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884\n",
-      "  warnings.warn(\n"
-     ]
-    }
-   ],
-   "source": [
-    "model = AutoModelForCausalLMWithValueHead.from_pretrained(\"gpt2\")\n",
-    "ref_model = AutoModelForCausalLMWithValueHead.from_pretrained(\"gpt2\")\n",
-    "tokenizer = GPT2Tokenizer.from_pretrained(\"gpt2\")\n",
-    "tokenizer.pad_token = tokenizer.eos_token"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 48,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from trl_rlhf_data import runner, ScriptArguments\n",
-    "import re\n",
-    "from dataclasses import dataclass\n",
-    "from typing import Dict, List, Optional\n",
-    "\n",
-    "from datasets import load_dataset\n",
-    "from transformers import HfArgumentParser"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 49,
-   "metadata": {},
-   "outputs": [
-    {
-     "ename": "TypeError",
-     "evalue": "runner() takes 0 positional arguments but 1 was given",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
-      "Cell \u001b[0;32mIn[49], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m dataset \u001b[38;5;241m=\u001b[39m \u001b[43mrunner\u001b[49m\u001b[43m(\u001b[49m\u001b[43mScriptArguments\u001b[49m\u001b[43m)\u001b[49m\n",
-      "\u001b[0;31mTypeError\u001b[0m: runner() takes 0 positional arguments but 1 was given"
-     ]
-    }
-   ],
-   "source": [
-    "dataset = runner(ScriptArguments)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  }
- ],
- "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.10.13"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 2
-}

ml/kto_dataset_processor.py

@@ -1,65 +1,210 @@
-from datasets import load_dataset, Dataset
+from datasets import Dataset, load_dataset
 import pandas as pd
-from pdb import set_trace as st
+from sklearn.model_selection import train_test_split
+import json
+from ipdb import set_trace as st
+from transformers import AutoTokenizer
+from enum import Enum
 
+class SupportedLanguages(str, Enum):
+    """Enumeration of supported languages"""
+    ENGLISH = "English"
+    DUTCH = "Dutch"
+    ITALIAN = "Italian"
+    SPANISH = "Spanish"
+    FRENCH = "French"
+    GERMAN = "German"
+    PORTUGUESE = "Portuguese"
+    RUSSIAN = "Russian"
+    CHINESE = "Chinese"
+    JAPANESE = "Japanese"
+    KOREAN = "Korean"
 
-def process_dataset_ultrafeedback():
+def transform_conversation(
+    entry: dict,
+    model_name: str,
+    max_history_turns: int = 10,
+    max_history_tokens: int = 4000
+) -> list:
+    """Transform conversation into KTO format with history"""
+    data_points = []
+    conversation = entry["conversation"]
+    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+
+    for i, message in enumerate(conversation):
+        # Only create data points for assistant messages that have ratings
+        if message["role"] != "assistant" or message["rating"] not in [1, -1]:
+            continue
+
+        # Get previous messages up to limits
+        formatted_history = []
+        formatted_prompt = ""
+        tokens = 0
+        pairs = 0  # Count complete user/assistant pairs
+
+        # Start from the current message and work backwards
+        current_idx = i - 1
+        while current_idx >= 0 and pairs < max_history_turns:
+            # We need both user and assistant messages to form a pair
+            if current_idx > 0 and conversation[current_idx]["role"] == "user" and conversation[current_idx-1]["role"] == "assistant":
+                # Add the pair to history
+                formatted_history.insert(0, conversation[current_idx-1])  # assistant
+                formatted_history.insert(1, conversation[current_idx])    # user
+
+                # Check token limit
+                try:
+                    current_formatted = tokenizer.apply_chat_template(formatted_history, tokenize=False)
+                    current_tokens = len(tokenizer.encode(current_formatted))
+
+                    if current_tokens > max_history_tokens:
+                        formatted_history = formatted_history[2:]  # Remove the oldest pair
+                        break
+
+                    formatted_prompt = current_formatted
+                    tokens = current_tokens
+                    pairs += 1
+                    current_idx -= 2
+                except Exception:
+                    # If template application fails, remove the last added pair
+                    formatted_history = formatted_history[2:]
+                    break
+            else:
+                current_idx -= 1
+
+        # Add the final user message that prompted the rated response
+        if i > 0 and conversation[i-1]["role"] == "user":
+            last_history = formatted_history + [conversation[i-1]]
+            try:
+                formatted_prompt = tokenizer.apply_chat_template(last_history, tokenize=False)
+            except Exception:
+                # If template application fails, use the previous valid prompt
+                pass
+
+        data_points.append({
+            "prompt": formatted_prompt.strip(),
+            "completion": message["content"].strip(),
+            "label": message["rating"] == 1,
+            "timestamp": entry["timestamp"],
+            "session_id": entry["session_id"],
+            "conversation_id": entry["conversation_id"],
+            "language": entry["language"]
+        })
+
+    return data_points
+
+def process_feel_dataset(
+    language: str,
+    model_name: str = "CohereForAI/aya-expanse-8b",
+    max_history_turns: int = 10,
+    max_history_tokens: int = 4000
+):
     """
-    Processes the 'train_prefs' and 'test_prefs' splits of the 'HuggingFaceH4/ultrafeedback_binarized' dataset
-    into a unified format for preference modeling.
+    Processes the feel dataset into a format suitable for KTO training using TRL.
+
+    Args:
+        language: Language to filter the dataset for (must be one of SupportedLanguages)
+        model_name: Name of the model to format for
+        max_history_turns: Maximum number of previous turns to include in history
+        max_history_tokens: Maximum number of tokens allowed in history
 
     Returns:
-        dict: A dictionary containing the unified 'train' and 'test' splits of the dataset in the KTO format.
-              Each split is a Hugging Face Dataset object.
+        dict: A dictionary containing the 'train' and 'test' splits of the dataset in KTO format
+
+    Raises:
+        ValueError: If language is not provided or not in SupportedLanguages
     """
-    # Load the relevant splits of the dataset
-    dataset_name = "HuggingFaceH4/ultrafeedback_binarized"
-    train_prefs = load_dataset(dataset_name, split="train_prefs")
-    test_prefs = load_dataset(dataset_name, split="test_prefs")
-
-    # Function to transform a single example into the desired schema
-    def transform_data(example):
-        data_points = []
-        # Chosen completion
-        chosen_completion = example["chosen"][1]["content"]
-        if chosen_completion.strip():  # Check for non-empty completions
-            data_points.append({
-                "prompt": example["prompt"],
-                "completion": chosen_completion.strip(),
-                "label": True
-            })
-        # Rejected completion
-        rejected_completion = example["rejected"][1]["content"]
-        if rejected_completion.strip():  # Check for non-empty completions
-            data_points.append({
-                "prompt": example["prompt"],
-                "completion": rejected_completion.strip(),
-                "label": False
-            })
-        return data_points
-
-    # Process train and test splits
-    train_data = []
-    test_data = []
-
-    for example in train_prefs:
-        train_data.extend(transform_data(example))
-
-    for example in test_prefs:
-        test_data.extend(transform_data(example))
-
-    # Convert unified data to DataFrames
-    train_df = pd.DataFrame(train_data)
-    test_df = pd.DataFrame(test_data)
+    # Validate language
+    if not language:
+        raise ValueError("Language parameter is required")
+
+    try:
+        # Validate that it's a supported language
+        SupportedLanguages(language)
+    except ValueError:
+        supported_langs = "\n- ".join([lang.value for lang in SupportedLanguages])
+        raise ValueError(
+            f"Invalid language: '{language}'\n"
+            f"Supported languages are:\n- {supported_langs}"
+        )
+
+    # Load feel dataset from HuggingFace
+    feel_dataset = load_dataset("feel-fl/feel-feedback")["train"]
+
+    # Filter dataset by language
+    feel_dataset = feel_dataset.filter(lambda x: x["language"] == language)
 
+    if len(feel_dataset) == 0:
+        raise ValueError(f"No data found for language: {language}")
+
+    kto_data = []
+
+    # Process all conversations in the filtered dataset
+    for entry in feel_dataset:
+        kto_data.extend(transform_conversation(
+            entry,
+            model_name,
+            max_history_turns,
+            max_history_tokens
+        ))
+
+    if len(kto_data) == 0:
+        raise ValueError(f"No valid training examples found for language: {language}")
+
+    # Convert to DataFrame
+    kto_df = pd.DataFrame(kto_data)
+
+    # Split into train and test sets (70% train, 30% test)
+    train_df, test_df = train_test_split(kto_df, test_size=0.3, random_state=42)
+
+    # Reset index to remove '__index_level_0__'
+    train_df = train_df.reset_index(drop=True)
+    test_df = test_df.reset_index(drop=True)
 
     # Convert to Hugging Face Dataset
-    unified_train = Dataset.from_pandas(train_df)
-    unified_test = Dataset.from_pandas(test_df)
+    train_dataset = Dataset.from_pandas(train_df)
+    test_dataset = Dataset.from_pandas(test_df)
 
-    return {"train": unified_train, "test": unified_test}
+    print(f"Processed {len(kto_data)} examples for language: {language}")
+    print(f"Train set size: {len(train_dataset)}")
+    print(f"Test set size: {len(test_dataset)}")
 
+    return {"train": train_dataset, "test": test_dataset}
 
 if __name__ == "__main__":
-    kto_dataset = process_dataset_ultrafeedback()
-    st()
+    # Process the dataset
+    datasets = process_feel_dataset("English")
+
+    # Print distribution of positive/negative labels
+    train_labels = datasets['train']['label']
+    test_labels = datasets['test']['label']
+
+    print("\nLabel Distribution:")
+    print("Train set:")
+    print(f"Positive feedback: {sum(train_labels)}")
+    print(f"Negative feedback: {len(train_labels) - sum(train_labels)}")
+    print(f"Positive ratio: {sum(train_labels)/len(train_labels):.2%}")
+
+    print("\nTest set:")
+    print(f"Positive feedback: {sum(test_labels)}")
+    print(f"Negative feedback: {len(test_labels) - sum(test_labels)}")
+    print(f"Positive ratio: {sum(test_labels)/len(test_labels):.2%}")
+
+    # Load original FEEL dataset
+    feel_dataset = load_dataset("feel-fl/feel-feedback", split="train")
+
+    # Print one original conversation
+    print("\nOriginal conversation from FEEL dataset:")
+    print(json.dumps(feel_dataset[0], indent=2))
+
+    # Print sample entries from processed dataset
+    print("\nSample entries from processed KTO dataset:")
+    print("\n" + "="*80 + "\nTRAIN SET SAMPLES\n" + "="*80)
+
+    # Export datasets to CSV
+    train_df = datasets['train'].to_pandas()
+    test_df = datasets['test'].to_pandas()
+
+    train_df.to_csv('kto_train_dataset.csv', index=False)
+    test_df.to_csv('kto_test_dataset.csv', index=False)
+
+    print("\nDatasets exported to 'kto_train_dataset.csv' and 'kto_test_dataset.csv'")

ml/kto_lora.py

@@ -0,0 +1,185 @@
+import os
+import torch
+from dataclasses import dataclass
+from accelerate import PartialState
+from transformers import AutoModelForCausalLM, AutoTokenizer, HfArgumentParser
+from trl import KTOConfig, KTOTrainer, ModelConfig, get_peft_config, maybe_unpair_preference_dataset, setup_chat_format
+from kto_dataset_processor import process_feel_dataset
+from datetime import datetime
+import wandb
+
+# PEFT library: attach and load adapters
+from peft import get_peft_model, PeftModel
+
+####################################
+#  CONFIGURATION
+####################################
+
+@dataclass
+class ScriptArguments:
+    """
+    Configuration for the script.
+    """
+    process_dataset_func: callable = process_feel_dataset  # Function to process dataset
+    checkpoint_path: str = None  # Checkpoint path if needed
+    push_to_hub: bool = False  # Whether to push the adapter to the HF Hub after training
+    language: str = "en"  # Language identifier (e.g., "en", "fr", etc.)
+
+@dataclass
+class ModelArguments(ModelConfig):
+    """
+    Configuration for the model.
+    """
+    model_name: str = "CohereForAI/aya-expanse-8b"
+    use_peft: bool = True
+    lora_target_modules: str = "all-linear"
+    lora_r: int = 16
+    lora_alpha: int = 16
+    trust_remote_code: bool = True
+
+@dataclass
+class TrainingArguments(KTOConfig):
+    """
+    Configuration for the KTO trainer.
+    """
+    output_dir: str = f"kto_{ModelArguments.model_name}_{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}"
+    num_train_epochs: int = 1
+    per_device_train_batch_size: int = 4
+    learning_rate: float = 5e-7
+    lr_scheduler_type: str = "cosine"
+    gradient_accumulation_steps: int = 1
+    logging_steps: int = 10
+    eval_steps: int = 500
+    warmup_ratio: float = 0.1
+    bf16: bool = True
+    logging_first_step: bool = True
+
+# Initialize configurations
+script_args = ScriptArguments()
+training_args = TrainingArguments()
+model_args = ModelArguments()
+
+####################################
+#  HELPER FUNCTIONS
+####################################
+
+def load_model_and_tokenizer(model_args):
+    """
+    Load the base model and tokenizer from the Hugging Face Hub.
+    """
+    model = AutoModelForCausalLM.from_pretrained(
+        model_args.model_name,
+        trust_remote_code=model_args.trust_remote_code,
+        torch_dtype=torch.float16,
+        device_map="auto"
+    )
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_args.model_name,
+        trust_remote_code=model_args.trust_remote_code
+    )
+
+    # Set pad token if it is missing
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    # Setup chat format if not available on the tokenizer
+    if not getattr(tokenizer, "chat_template", None):
+        model, tokenizer = setup_chat_format(model, tokenizer)
+
+    return model, tokenizer
+
+####################################
+#  MAIN LOGIC
+####################################
+
+def main():
+    # Initialize wandb for logging
+    wandb.init(project="kto")
+
+    print("Loading base model and tokenizer...")
+    model, tokenizer = load_model_and_tokenizer(model_args)
+    ref_model, _ = load_model_and_tokenizer(model_args)
+    print("Models and tokenizer loaded.")
+
+    # -----------------------------
+    # Adapter Loading or Initialization
+    # -----------------------------
+    # Configure the PEFT / LoRA adapter settings
+    peft_config = get_peft_config(model_args)
+    adapter_dir = os.path.join("adapters", script_args.language)
+
+    if os.path.isdir(adapter_dir):
+        # If an adapter for this language already exists, load it into the base model.
+        model = PeftModel.from_pretrained(model, adapter_dir)
+        print(f"Loaded existing adapter for language '{script_args.language}' from {adapter_dir}.")
+    else:
+        # Otherwise, initialize a new LoRA adapter.
+        model = get_peft_model(model, peft_config)
+        print(f"No adapter found for language '{script_args.language}'. Initialized new adapter.")
+
+    # -----------------------------
+    # Data Preparation and Training
+    # -----------------------------
+    print("Processing dataset...")
+    dataset = script_args.process_dataset_func()
+    print("Dataset processed.")
+
+    print("Initializing trainer...")
+    trainer = KTOTrainer(
+        model=model,
+        ref_model=ref_model,
+        args=training_args,
+        train_dataset=dataset["train"],
+        eval_dataset=dataset["test"],
+        processing_class=tokenizer,
+        peft_config=peft_config,
+    )
+
+    # Training
+    print("Starting training...")
+    trainer.train()
+    print("Training completed.")
+
+    # Evaluation
+    print("Evaluating model...")
+    metrics = trainer.evaluate()
+    print(f"Metrics: {metrics}")
+    trainer.log_metrics("eval", metrics)
+    trainer.save_metrics("eval", metrics)
+
+    # Log metrics to wandb
+    wandb.log({
+        "epoch": metrics.get("epoch"),
+        "grad_norm": metrics.get("grad_norm"),
+        "kl": metrics.get("kl"),
+        "learning_rate": metrics.get("learning_rate"),
+        "logits/chosen": metrics.get("logits/chosen"),
+        "logits/rejected": metrics.get("logits/rejected"),
+        "logps/chosen": metrics.get("logps/chosen"),
+        "logps/rejected": metrics.get("logps/rejected"),
+        "loss": metrics.get("loss"),
+        "rewards/chosen": metrics.get("rewards/chosen"),
+        "rewards/margins": metrics.get("rewards/margins"),
+        "rewards/rejected": metrics.get("rewards/rejected"),
+        "step": metrics.get("step")
+    })
+
+    # -----------------------------
+    # Adapter Saving
+    # -----------------------------
+    print("Saving adapter...")
+    os.makedirs(adapter_dir, exist_ok=True)
+    model.save_pretrained(adapter_dir)
+    print(f"Adapter for language '{script_args.language}' saved to: {adapter_dir}")
+
+    if script_args.push_to_hub:
+        print("Pushing adapter to Hugging Face Hub...")
+        model.push_to_hub(repo_id=f"your_hf_org/{script_args.language}-adapter")
+
+    print("Process completed.")
+
+    # Finish wandb run
+    wandb.finish()
+
+if __name__ == "__main__":
+    main()