{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "68f38794",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🖥️  Using device: cuda\n",
      "GPU: NVIDIA GeForce RTX 4060 Ti\n",
      "Memory: 15.6 GB\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "# Check GPU and Set Device\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "print(f\"🖥️  Using device: {device}\")\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    print(f\"GPU: {torch.cuda.get_device_name(0)}\")\n",
    "    print(f\"Memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB\")\n",
    "else:\n",
    "    print(\"⚠️  No GPU available, using CPU (training will be slower)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9ac8e7db",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/miriam/miniconda3/envs/case-dev/lib/python3.10/site-packages/requests/__init__.py:86: RequestsDependencyWarning: Unable to find acceptable character detection dependency (chardet or charset_normalizer).\n",
      "  warnings.warn(\n",
      "/home/miriam/miniconda3/envs/case-dev/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "from transformers import T5ForConditionalGeneration, T5Tokenizer\n",
    "\n",
    "def load_trained_model(model_path):\n",
    "    \"\"\"Load a trained model\"\"\"\n",
    "    model = T5ForConditionalGeneration.from_pretrained(model_path)\n",
    "    tokenizer = T5Tokenizer.from_pretrained(model_path)\n",
    "    model.to(device)\n",
    "    print(f\"✅ Model loaded from: {model_path}\")\n",
    "    return model, tokenizer\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d66a9e03",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ Model loaded from: ./output/flan_t5_case_summarization/best_model\n"
     ]
    }
   ],
   "source": [
    "model_path = \"./output/flan_t5_case_summarization/best_model\"\n",
    "model, tokenizer = load_trained_model(model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "5b3f85fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configuration\n",
    "config = {\n",
    "    'model_name': 'google/flan-t5-base',\n",
    "    'train_data_path': 'data/train_data.jsonl',  # Change to your data path\n",
    "    'val_data_path': 'data/test_data.jsonl',      # Change to your data path\n",
    "    'output_dir': 'output/flan_t5_case_summarization',\n",
    "    'batch_size': 4,                    # Adjust based on your GPU memory\n",
    "    'learning_rate': 3e-5,\n",
    "    'weight_decay': 0.01,\n",
    "    'adam_epsilon': 1e-8,\n",
    "    'num_epochs': 3,\n",
    "    'max_source_length': 1024,\n",
    "    'max_target_length': 256,\n",
    "    'warmup_ratio': 0.1,\n",
    "    'save_checkpoints': True,\n",
    "    'use_wandb': False,                 # Set to True if you want to use wandb\n",
    "    'wandb_project': 'flan-t5-case-summarization',\n",
    "    'run_name': 'flan-t5-large-legal-cases'\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d5d35bf8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_summary(text: str, max_length: int = 256) -> str:\n",
    "    \"\"\"Generate summary of a child helpline case call transcript\"\"\"\n",
    "    input_text = f\"Summarize the following child helpline case call transcript:{text}\"\n",
    "    \n",
    "    inputs = tokenizer(\n",
    "        input_text,\n",
    "        max_length=config['max_source_length'],\n",
    "        padding='max_length',\n",
    "        truncation=True,\n",
    "        return_tensors='pt'\n",
    "    ).to(device)\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        outputs = model.generate(\n",
    "            input_ids=inputs['input_ids'],\n",
    "            attention_mask=inputs['attention_mask'],\n",
    "            max_length=max_length,\n",
    "            num_beams=4,\n",
    "            length_penalty=2.0,\n",
    "            early_stopping=True,\n",
    "            no_repeat_ngram_size=2\n",
    "        )\n",
    "    \n",
    "    summary = tokenizer.decode(outputs[0], skip_special_tokens=True)\n",
    "    return summary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9359048d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "John reported a case of child labor involving his 12-year-old sister in Mwanza. The counselor advised him to report it to the local Labor Office and police, with follow-up from the helpline.\n"
     ]
    }
   ],
   "source": [
    "# Test\n",
    "text = \"Hi, is this 116? Yes, thank you for calling. Who am I speaking to? My name is John, I'm from Mwanza. I've got a serious concern about my 12-year-old sister. She's being forced into child labor at a local factory. This sounds terrible, John. Thank you for your bravery. Have others noticed this? Yes, some workers have expressed their concerns but are too afraid to speak up. She's always tired and has bruises. I'm scared too. I understand. The best thing you can do is report it to the Mwanza Labor Office and also to the police. We can follow up on your call as well. Please don't hesitate to call again if there are any updates.\"\n",
    "summary = generate_summary(text)\n",
    "print(summary)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7558272e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mariam reported a case of child marriage involving her dad in Kisauni. The counselor advised reporting the issue to the children's office and police, and offered follow-up support.\n"
     ]
    }
   ],
   "source": [
    "# Test\n",
    "text = \"Hello, is this 116? Yes, thank you for calling. Who am I speaking to? You're speaking to Mariam from Kisauni. My dad wants to forcefully marry me off to his old friend. The wedding is set to take place tomorrow at 10 am. I am so sorry mariam. Thank you for reporting this for it is wrong and against the law . Is your mum aware? Yes, but are too afraid to speak up. She's always tired and has bruises. I'm scared too. I understand. The best thing you can do is report it to the Kisauni children's office and also to the police. We can follow up on your call as well. Please don't hesitate to call again if there are any updates.\"\n",
    "summary = generate_summary(text)\n",
    "print(summary)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "87432592-17a4-403a-8ec3-0255d2f7bfaa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Caller Felice reported a case of child labor involving his 5-month-old son. The child is still in the womb and is allowed to drink his mother's milk, but the mother has to breastfeed him for six months.\n"
     ]
    }
   ],
   "source": [
    "# Test\n",
    "text = \"\"\"\n",
    "Helpline call: Hello how are you  I'm good I'm  Felice What can I do for you I would like to ask you for a favor Yes The child is still in the womb How many days before the baby is to be born I don't know when the baby is born here Yes He is allowed to drink his mother's milk His big meal is his mother's milk Thank you What we want to do is for the mother to get food in time good food healthy food to help her to get better to get strength and to be able to prepare food for the baby The baby's food is taken from the mother's food Okay And he has to be breastfed for six months That's the months She has to breastfeed from the day he was born until the fifth month\n",
    "Task: Summarize the caller's request and the health advice provided.\n",
    "\"\"\"\n",
    "summary = generate_summary(text)\n",
    "print(summary)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4aa1163-7005-4dcd-9188-f6a57ada85c9",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": " case-dev ",
   "language": "python",
   "name": "case-dev"
  },
  "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.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}