{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from datasets import load_dataset\n",
    "from transformers import GPT2Tokenizer, GPT2LMHeadModel, TrainingArguments, Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the DailyDialog dataset\n",
    "dataset = load_dataset('daily_dialog')\n",
    "\n",
    "# Concatenate all utterances within a dialogue and map to 'dialog' key\n",
    "def concatenate_utterances(example):\n",
    "    example['dialog'] = \" \".join(example['dialog'])\n",
    "    return example\n",
    "\n",
    "# Apply the function to all examples in the dataset\n",
    "dataset = dataset.map(concatenate_utterances)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the tokenizer and model\n",
    "tokenizer = GPT2Tokenizer.from_pretrained('microsoft/DialoGPT-medium')\n",
    "tokenizer.pad_token = tokenizer.eos_token\n",
    "model = GPT2LMHeadModel.from_pretrained('microsoft/DialoGPT-medium')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5d576321ac974a118f75b83cd8437256",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/1000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8e7254605abe41dbad8d6b2321d904c5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/1000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Encode the dataset\n",
    "def encode(examples):\n",
    "    encoded = tokenizer(examples['dialog'], truncation=True, padding='max_length', max_length=128)\n",
    "    encoded['labels'] = encoded['input_ids'][:]\n",
    "    return encoded\n",
    "\n",
    "encoded_dataset = dataset.map(encode, batched=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define training arguments\n",
    "training_args = TrainingArguments(\n",
    "    output_dir='model',              # output directory\n",
    "    num_train_epochs=2,              # total number of training epochs\n",
    "    per_device_train_batch_size=64,  # batch size per device during training\n",
    "    per_device_eval_batch_size=64,   # batch size for evaluation\n",
    "    warmup_steps=500,                # number of warmup steps for learning rate scheduler\n",
    "    weight_decay=0.01,               # strength of weight decay\n",
    "    logging_dir=None,                # directory for storing logs\n",
    ")\n",
    "\n",
    "# Create Trainer\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=encoded_dataset['train'],\n",
    "    eval_dataset=encoded_dataset['validation']\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2a46f1165bba4ef8b597eace733e9eaf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/16 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1214b0347e7149c88d02d31dbc53f523",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Evaluate before fine-tuning\n",
    "pre_eval_results = trainer.evaluate(encoded_dataset['validation'])\n",
    "\n",
    "# Get predictions for validation set before fine tuning for 10 samples\n",
    "pre_val_predictions = trainer.predict(encoded_dataset['validation'].select(range(10)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3f2c5fb0a67449a39c76d842dc4e6ed5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/348 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'train_runtime': 25354.0984, 'train_samples_per_second': 0.877, 'train_steps_per_second': 0.014, 'train_loss': 2.2603482651984557, 'epoch': 2.0}\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=348, training_loss=2.2603482651984557, metrics={'train_runtime': 25354.0984, 'train_samples_per_second': 0.877, 'train_steps_per_second': 0.014, 'train_loss': 2.2603482651984557, 'epoch': 2.0})"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Fine-tune the model\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6a3612bcb25b4efcbb39d2c15d048b07",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0e147b6f77194fb58a7f1869dbd46be0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/16 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Evaluation Results before fine-tuning : 4.766543388366699\n",
      "Evaluation Results after fine-tuning  : 1.8690917491912842\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fbc53f3e0f2348928c5321f296dfb472",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Get predictions for validation set before fine tuning for 10 samples\n",
    "pre_val_predictions = trainer.predict(encoded_dataset['validation'].select(range(10)))\n",
    "\n",
    "# Evaluate after fine-tuning\n",
    "post_eval_results = trainer.evaluate(encoded_dataset['validation'])\n",
    "\n",
    "# Print the evaluation losses before and after fine-tuning\n",
    "print('Evaluation Results before fine-tuning :', pre_eval_results['eval_loss'])\n",
    "print('Evaluation Results after fine-tuning  :', post_eval_results['eval_loss'])\n",
    "\n",
    "# Get predictions for validation set before fine tuning for 10 samples\n",
    "post_val_predictions = trainer.predict(encoded_dataset['validation'].select(range(10)))\n",
    "\n",
    "# Zip the pre and post tuning predictions\n",
    "predictions = zip(pre_val_predictions.predictions, post_val_predictions.predictions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ground truth \n",
      "Good morning , sir . Is there a bank near here ?   There is one . 5 blocks away from here ?   Well , that's too far.Can you change some money for me ?   Surely , of course . What kind of currency have you got ?   RIB .   How much would you like to change ?   1000 Yuan.Here you are . \n",
      "\n",
      "Pre-prediction \n",
      " and, sir.  there anything problem here here?   Yes is. in, away. here.   Yes, I's a far.How you tell the money for me?   Sure. sir course. Here's of money do you got?   IIB.   Here much is you like to exchange?   I R.How you are.   \n",
      "\n",
      "Post-prediction \n",
      " and, sir.  there anything problem here here?   Yes is. in, away. here.   Yes, I's a far.How you tell the money for me?   Sure. sir course. Here's of money do you got?   IIB.   Here much is you like to exchange?   I R.How you are.   \n",
      "\n",
      "----------------------------------------------------------------------------------------------------------------------\n",
      "\n",
      "Ground truth \n",
      "Good afternoon . This is Michelle Li speaking , calling on behalf of IBA . Is Mr Meng available at all ?   This is Mr Meng speaking , Michelle .   Oh , hello ! Sorry about that . I'm just calling to say that we've received your new Corporate Credit Card from HQ .   That was quick ! I wasn't expecting it until later this week .   Yes , our application procedures have speeded up since we started using the new fast-track system .   Shall I come in and collect it ?   Or we can send it to you . But if you would like to use it at the ATM , you'll need to wait for your PIN number .   Mmmm ... if I come in and collect it this afternoon , is there any way I could use it today ? Petty cash is getting low , so I need to draw some money .   As long as you bring your ID , etc , we can serve you over the counter . But you won't be able to use the ATM until your new PIN number arrives .   I see . Yes , that's fine . I'll be there at around 2:30 pm . See you later , and thanks . \n",
      "\n",
      "Pre-prediction \n",
      " and, I is Mr... the for you of theKE.  there. here? the?   Yes is Mr Meng.. calling Li I  I, I. I, the. I'm sorry a to see that I have got your letter contract Account Card. China.   Oh's fast. Thank'm't expecting that to now. afternoon.   I, I new was are beened up. then received working it new card cardtrack system.   I we call in now check it?   Yes you can go it to you.  you you don like to come it,\n",
      "\n",
      "Post-prediction \n",
      " and, I is Mr... the for you of theKE.  there. here? the?   Yes is Mr Meng.. calling Li I  I, I. I, the. I'm sorry a to see that I have got your letter contract Account Card. China.   Oh's fast. Thank'm't expecting that to now. afternoon.   I, I new was are beened up. then received working it new card cardtrack system.   I we call in now check it?   Yes you can go it to you.  you you don like to come it,\n",
      "\n",
      "----------------------------------------------------------------------------------------------------------------------\n",
      "\n",
      "Ground truth \n",
      "What qualifications should a reporter have ?   As a reporter , he must have acute insight and language skills . At the same time , he must have good judgment , the respect for his job and tactical cooperation with others .   Can you work under pressure ? You know , people working here are all busy everyday since we're daily newspaper .   I think I've got used to work under pressure . I will adjust myself to the step of your newspaper quickly . \n",
      "\n",
      "Pre-prediction \n",
      " and do I person be?   I a reporter, I should be a knowledge and knowledge skills.  least same time, he must be a knowledge. and ability of the colleagues, the ability. the.   What you tell for pressure?  must, the are under are always very.. the have in news.   Yes know so can got the to it under pressure. I'm try my. the new of pressure question..  \n",
      "\n",
      "Post-prediction \n",
      " and do I person be?   I a reporter, I should be a knowledge and knowledge skills.  least same time, he must be a knowledge. and ability of the colleagues, the ability. the.   What you tell for pressure?  must, the are under are always very.. the have in news.   Yes know so can got the to it under pressure. I'm try my. the new of pressure question..  \n",
      "\n",
      "----------------------------------------------------------------------------------------------------------------------\n",
      "\n",
      "Ground truth \n",
      "Hi , good morning , Miss ? what can I help you with ?   Good morning I'd like to mail this box of books to Taiwan .   OK , please put it on this scale.Airmail or by sea ?   How long does it take to reach Taiwan by sea ?   Usually about two month .   That's too long.How long does it take to reach Taiwan by airmail ?   About ten days .   Then how much is that by airmail ?   Let me see.It ' s 57 dollars , 20 cents , including tax .   That's a little bit expensive .   Although it's expensive to send it by airmail , it's quicker and safer than by sea .   I guess I have to send it by airmail .   Do you want to ensure the contents , Miss ?   Yes , please .   Please fill out this form , also please write the value of the items in this space .   OK . \n",
      "\n",
      "Pre-prediction \n",
      " and I morning. I.  can I do you with?   I morning,'m like to buy some letter to envelop to my.   I, I send it in the shelf.How fewail. express air.   By about will it take to ship China? sea?   About it ten weeks.   OK's not long.How much does it take to get China by air??   About three days.   That I long is it? seamail?   About me see.It's s about yuan. including dollars per and postage.   That's fine little \n",
      "\n",
      "Post-prediction \n",
      " and I morning. I.  can I do you with?   I morning,'m like to buy some letter to envelop to my.   I, I send it in the shelf.How fewail. express air.   By about will it take to ship China? sea?   About it ten weeks.   OK's not long.How much does it take to get China by air??   About three days.   That I long is it? seamail?   About me see.It's s about yuan. including dollars per and postage.   That's fine little \n",
      "\n",
      "----------------------------------------------------------------------------------------------------------------------\n",
      "\n",
      "Ground truth \n",
      "Excuse me , ma'am . Can you tell me where the nearest postoffice is ?   Of course . Go straight ahead . Turn right at the next street . You'll see a tall , yellow building.The post office is on the first floor .   Do you mean that I go that way for one block , then turn right ?   Yes , you are right .   Is it far ?   No , It's only about five minutes ' walk .   Thank you very much .   It's my pleasure . \n",
      "\n",
      "Pre-prediction \n",
      " and me, sir'am. I I tell me where the bus train office is?   Yes course. It to to.  right at the first intersection. 'll see it post building white building.It post office is on the right floor.   Thank you know the the can straight way? the post? or turn right at   Yes. that go right. The  Thank there a from   No, it's not a a blocks walk walk.   I you. much.  You's a pleasure.  \n",
      "\n",
      "Post-prediction \n",
      " and me, sir'am. I I tell me where the bus train office is?   Yes course. It to to.  right at the first intersection. 'll see it post building white building.It post office is on the right floor.   Thank you know the the can straight way? the post? or turn right at   Yes. that go right. The  Thank there a from   No, it's not a a blocks walk walk.   I you. much.  You's a pleasure.  \n",
      "\n",
      "----------------------------------------------------------------------------------------------------------------------\n",
      "\n",
      "Ground truth \n",
      "Could you give me some advice on how to bring up my son properly ?   He's a bright boy , isn't he ?   But he always wimps out of difficulty .   Don't worry , he'll make good progress step by step . \n",
      "\n",
      "Pre-prediction \n",
      " and tell me a advice? how to get my my resume??   You's a good young. isn't he?   Yes he's getsagsps on of bed.   I't worry, he'll get up use. by step.   \n",
      "\n",
      "Post-prediction \n",
      " and tell me a advice? how to get my my resume??   You's a good young. isn't he?   Yes he's getsagsps on of bed.   I't worry, he'll get up use. by step.   \n",
      "\n",
      "----------------------------------------------------------------------------------------------------------------------\n",
      "\n",
      "Ground truth \n",
      "I'm in 507 . I have a few problems with my room .   What is that problem , sir ?   There are cockroaches in my room .   Are you sure , sir ? Flies I could believe , but cockroaches ?   I've counted nine different cockroaches , and I accidentally stepped on another one .   Sir , we run a spotless and cockroach-less hotel .   You dare to doubt me ?   I'm sorry , sir . Let me transfer you to my supervisor . \n",
      "\n",
      "Pre-prediction \n",
      " and sorry the8.'m a friend questions. my car.   What are the?? Mr?   I's aroaches in my room.   What you sure you sir? ies are see see. but cockroaches?   Yes'm never them cock cockroaches in sir I've counted on one one.   I, I have a bug check room cleanroach freefree room.   What mean to step the?   I'm not, sir. I me check you to the room.   \n",
      "\n",
      "Post-prediction \n",
      " and sorry the8.'m a friend questions. my car.   What are the?? Mr?   I's aroaches in my room.   What you sure you sir? ies are see see. but cockroaches?   Yes'm never them cock cockroaches in sir I've counted on one one.   I, I have a bug check room cleanroach freefree room.   What mean to step the?   I'm not, sir. I me check you to the room.   \n",
      "\n",
      "----------------------------------------------------------------------------------------------------------------------\n",
      "\n",
      "Ground truth \n",
      "Excuse me , sir , I'm afraid you can't park your car here .   Why not ? It's my parking space .   I'm afraid not , sir .   Oh ? That's a surprise . Let me see ... D 0411 Our dog's birthday . Yes , I'm sure this my parking space !   But I saw a red car always parking here before .   Oh , we've just repainted our car . It was red .   Maybe . But the car of this space has a broken rearview mirror on the left .   Yeah . It used to . We got that fixed yesterday too .   Could you wait for a minute , sir ? I'd like to have a check .   Sure , go ahead .   Sorry , sir , my mistake . This is your parking space .   That's all right . It's not your fault . \n",
      "\n",
      "Pre-prediction \n",
      " and me, sir. can'm afraid I have't have here car here.   I not?'s a car space.   I'm sorry I. sir.   I, What's a shame.  me see. ang., car is name is , it'm afraid you is car space.   Oh I'm a car car parked parked there..   I, I have got gotainted it car.'s a before   Oh you But it car is the car is a red window window mirror. the left side   Oh, I's to be've it\n",
      "\n",
      "Post-prediction \n",
      " and me, sir. can'm afraid I have't have here car here.   I not?'s a car space.   I'm sorry I. sir.   I, What's a shame.  me see. ang., car is name is , it'm afraid you is car space.   Oh I'm a car car parked parked there..   I, I have got gotainted it car.'s a before   Oh you But it car is the car is a red window window mirror. the left side   Oh, I's to be've it\n",
      "\n",
      "----------------------------------------------------------------------------------------------------------------------\n",
      "\n",
      "Ground truth \n",
      "What can I do for you today ?   I need to buy a new refrigerator today .   Were you looking at a particular refrigerator ?   I like that Kenmore refrigerator .   This particular refrigerator is a very good choice .   Tell me about it .   Not only is it affordable , but it comes with all the appliances .   What are the appliances .   It has an ice maker , water dispenser , and plenty of room on the inside .   I'd like to see it for myself .   Go right ahead .   I like what I see . \n",
      "\n",
      "Pre-prediction \n",
      " and I do for you??   I'd to get a new car..   What you able at the new brand?   Yes was the onemore refrigerator.   What one refrigerator is a Ken popular refrigerator.   I me more it.   It only is it a, it it's with a the appliances.   I appliances the appliances?   They comes a electric maker, a heaterer, and a of other for the bottom.   What like like to buy it. a.   It ahead ahead.   I'll it I see.   \n",
      "\n",
      "Post-prediction \n",
      " and I do for you??   I'd to get a new car..   What you able at the new brand?   Yes was the onemore refrigerator.   What one refrigerator is a Ken popular refrigerator.   I me more it.   It only is it a, it it's with a the appliances.   I appliances the appliances?   They comes a electric maker, a heaterer, and a of other for the bottom.   What like like to buy it. a.   It ahead ahead.   I'll it I see.   \n",
      "\n",
      "----------------------------------------------------------------------------------------------------------------------\n",
      "\n",
      "Ground truth \n",
      "Oh , well . It was fun to be the winner . But ... it's too big . I must be an extra small in the States .   So what about the tennis racket ?   Look ! It's amazing . I can't wait to try it out !   How much did that end up costing you ?   Oh ... around twenty bucks . A bargain if you ask me . Look at the picture of her playing with it !   Hey , two for one . That's a super deal .   And here's her signature ! \n",
      "\n",
      "Pre-prediction \n",
      " and I, I's nice. meet there first.  I I was not late. I'm go the old.. the company.   You,? the other court?   It, It's a! It can't believe to use it..   I much is you cost up costing??   It, it ten dollars.  little! you ask me. , the racket. the.. it.   Oh, I dollars twenty!'s a bargain bargain!   I you I another purse.   \n",
      "\n",
      "Post-prediction \n",
      " and I, I's nice. meet there first.  I I was not late. I'm go the old.. the company.   You,? the other court?   It, It's a! It can't believe to use it..   I much is you cost up costing??   It, it ten dollars.  little! you ask me. , the racket. the.. it.   Oh, I dollars twenty!'s a bargain bargain!   I you I another purse.   \n",
      "\n",
      "----------------------------------------------------------------------------------------------------------------------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for idx, (pre, post) in enumerate(predictions):\n",
    "    pre_pred = tokenizer.decode(np.argmax(pre, axis=-1), skip_special_tokens=True)\n",
    "    post_pred = tokenizer.decode(np.argmax(post, axis=-1), skip_special_tokens=True)\n",
    "    ground_truth = encoded_dataset['validation'][idx][\"dialog\"]\n",
    "    \n",
    "    print('Ground truth \\n' + ground_truth + '\\n')\n",
    "    print('Pre-prediction \\n' + pre_pred + '\\n')\n",
    "    print('Post-prediction \\n'+ post_pred + '\\n')\n",
    "    print('----------------------------------------------------------------------------------------------------------------------\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer.save_pretrained(\"saved_model\")\n",
    "model.save_pretrained(\"saved_model\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.11.5"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}