kraken_train_router.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d13fa7d8-4b83-4fd8-80aa-8f064a52997b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install accelerate transformers datasets scikit-learn flash-attn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb3c80be-b021-4881-b8ef-5c86d6a6164a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36525529-8c42-4aae-b877-89c33abe6e30",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, Trainer, DataCollatorWithPadding\n",
    "from datasets import Dataset\n",
    "from sklearn.metrics import classification_report\n",
    "from sklearn.model_selection import train_test_split\n",
    "import pandas as pd\n",
    "import glob\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fcc1c30d-9964-48eb-a7bd-463b927eb1d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pfad zu deinen Parquet-Dateien\n",
    "parquet_files_path = \"./your_router_data_path/*.parquet\"\n",
    "\n",
    "# Alle Parquet-Dateien laden und konkatinieren\n",
    "parquet_files = glob.glob(parquet_files_path)\n",
    "data_frames = [pd.read_parquet(file) for file in parquet_files]\n",
    "data = pd.concat(data_frames, ignore_index=True)\n",
    "\n",
    "# Überprüfen, ob die Spaltennamen korrekt sind\n",
    "print(data.columns)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8cf23bf-3332-413f-8f19-0805fa8ca830",
   "metadata": {},
   "outputs": [],
   "source": [
    "split = 0.90\n",
    "RANDOM_SEED = 42\n",
    "\n",
    "df_train, df_test_unfixed = train_test_split(data, test_size=(1-split), random_state=RANDOM_SEED)\n",
    "df_val, df_test = train_test_split(df_test_unfixed, test_size=0.5, random_state=RANDOM_SEED)\n",
    "\n",
    "df_train = df_train.reset_index(drop=True)\n",
    "df_test_unfixed = df_test_unfixed.reset_index(drop=True)\n",
    "df_val = df_val.reset_index(drop=True)\n",
    "df_test = df_test.reset_index(drop=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "262eb005-a187-41ee-bc20-7b5f68cc0231",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load tokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"Qwen/Qwen1.5-0.5B\")\n",
    "tokenizer.pad_token = tokenizer.eos_token\n",
    "\n",
    "# Load sequence classification model and adjust head\n",
    "model = AutoModelForSequenceClassification.from_pretrained(\"Qwen/Qwen1.5-0.5B\", device_map=\"auto\", num_labels=5)\n",
    "model.config.pad_token_id = tokenizer.pad_token_id\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7713ee83-b74e-4188-a398-efd72c7c0fb9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess_function(examples):\n",
    "    return tokenizer(examples[\"text\"], truncation=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "114673b0-bb06-40da-af7d-2558dc2cc02f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prepare datasets for training\n",
    "ds_train = Dataset.from_pandas(df_train)\n",
    "ds_eval = Dataset.from_pandas(df_val)\n",
    "ds_test = Dataset.from_pandas(df_test)\n",
    "\n",
    "# Apply preprocessing to the datasets\n",
    "ds_train = ds_train.map(preprocess_function, batched=True)\n",
    "ds_eval = ds_eval.map(preprocess_function, batched=True)\n",
    "ds_test = ds_test.map(preprocess_function, batched=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b3c0285-8f2a-48fa-829b-eaf360b9ef5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training Argumente\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=\"./select_your_router_path\",\n",
    "    learning_rate=2e-5,\n",
    "    per_device_train_batch_size=4,\n",
    "    per_device_eval_batch_size=4,\n",
    "    num_train_epochs=7,\n",
    "    weight_decay=0.1,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b488726-f0b0-4c4f-a66c-1cf9d2f1e9b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Data Collator\n",
    "data_collator = DataCollatorWithPadding(tokenizer=tokenizer)\n",
    "\n",
    "# Initialize trainer\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=ds_train,\n",
    "    eval_dataset=ds_eval,\n",
    "    tokenizer=tokenizer,\n",
    "    data_collator=data_collator,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "192ea8b3-01e2-48e4-8724-bf8a0e0cda4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training\n",
    "trainer.train()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "841d7c77-156d-4acc-b526-d6be65eb44db",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluation auf dem Testset\n",
    "preds = trainer.predict(ds_test)\n",
    "\n",
    "# Umwandlung der Vorhersagen und wahren Labels in die benötigte Form\n",
    "true_labels = ds_test[\"label\"]\n",
    "pred_labels = preds.predictions.argmax(-1)\n",
    "\n",
    "# Classification Report\n",
    "report = classification_report(true_labels, pred_labels)\n",
    "print(report)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d050111f-1a1d-4204-af80-65917ca8ae86",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluation on the training set\n",
    "preds_train = trainer.predict(ds_train)\n",
    "true_labels_train = ds_train[\"label\"]\n",
    "pred_labels_train = preds_train.predictions.argmax(-1)\n",
    "report_train = classification_report(true_labels_train, pred_labels_train)\n",
    "print(report_train)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db11c60c-e37d-4ef0-82fc-d31c683666b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "largest_score = []\n",
    "\n",
    "iter = 0\n",
    "\n",
    "for sample in ds_train:\n",
    "    iter = iter + 1\n",
    "    print(iter)\n",
    "    largest_score.append(pipe(sample[\"text\"])[0][\"score\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "991723ac-d4ec-4963-b236-f916798a8aa1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "np.mean(largest_score)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca2abc75-3b6b-4df1-a1de-64ecd9948838",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.save_model(f\"./select_your_router_path/router\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a00162e-4ada-426b-84df-c8d5dcb2adec",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, AutoModelForSequenceClassification, TextClassificationPipeline\n",
    "model_name = './router_lang/checkpoint-50500'\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "model = AutoModelForSequenceClassification.from_pretrained(model_name)\n",
    "classifier = TextClassificationPipeline(model=model, tokenizer=tokenizer)\n",
    "res = classifier(\"Test your router model here to see if it fits your requirements\")\n",
    "print(res)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}