DSPy_Optimization.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "8b3ee6e2-ca9c-40fa-b4c6-a9596f075f79",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-04-22T23:03:20.101831Z",
     "iopub.status.busy": "2025-04-22T23:03:20.101435Z",
     "iopub.status.idle": "2025-04-22T23:03:20.105088Z",
     "shell.execute_reply": "2025-04-22T23:03:20.104580Z",
     "shell.execute_reply.started": "2025-04-22T23:03:20.101804Z"
    }
   },
   "outputs": [],
   "source": [
    "import dspy\n",
    "from dspy.teleprompt import MIPROv2\n",
    "from typing import List, Dict\n",
    "import json\n",
    "import numpy as np\n",
    "import os\n",
    "import random\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "4ec9a29b-9162-4fe3-b32d-4de4397c6483",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-04-22T23:00:21.439753Z",
     "iopub.status.busy": "2025-04-22T23:00:21.439342Z",
     "iopub.status.idle": "2025-04-22T23:00:21.526091Z",
     "shell.execute_reply": "2025-04-22T23:00:21.525575Z",
     "shell.execute_reply.started": "2025-04-22T23:00:21.439727Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4/4 [00:00<00:00, 77.75it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'input': {'src_text': 'Ma io che ne so, comandà? Io stavo a casa di mia madre, lo sapete.\\n\\nLo so.',\n",
       "  'tgt_text': \"What do I know, Commander? I was at my mom's house, you know it.\\n\\nI knows.\",\n",
       "  'src_prev': \"Questa è una linea. Qua faccio quello che voglio, è terra mia, la legge è mia. Dall'altro lato c'è un mondo fatto di spazzatura. Questa linea non l'ho mai oltrepassata. Impara chi è tua madre una volta per tutte. Tieni, questo era per te. Mà… Mà! Secondo me non è stata lei. Come al solito ti sei fatto prendere per il culo. Comandà, credo che non è stata lei. Carmine, sei uno stronzo. Robè, portalo via. Andiamocene.\",\n",
       "  'src_next': 'E allora che altro vi devo dire? Tu non devi dire niente. Devi tenere la bocca chiusa. E non dire a nessuno quello che ti ho detto. Ma a nessuno però. Ho capito. Però devi tenere le orecchie aperte e ascoltare tutto quello che si dice qua dentro. Perché prima o poi, chi fa queste cose parla. Si deve atteggiare, si deve fare grosso. Che si è divertito con la moglie del comandante. Secondo me vi sbagliate, comandà. Non può essere stato nessuno che sta qua dentro. Lo so.',\n",
       "  'tgt_prev': \"This is a line. Here I do whatever I want, it's my territory, it's my law. On the other side there's a world of trash. I've never crossed that line. Learn who your mother is, once and for all. Here, this was for you. Ma… Ma! I don't think it was her. As usual you let her fuck you around. Commander, I think she didn't do it. Carmine, you're an asshole. Robè, take him away. Let's go.\",\n",
       "  'tgt_next': \"So what else should I say? You don't have to say anything. You have to keep your mouth shut. And don't tell anybody what I told you. To nobody. Got it. But keep your ears open and listen to what they say in here. Because sooner or later, guys who do such things talk. They need to swagger, act like big guys. Bragging they had fun with the Commander's wife. I think you're wrong, Commander. It can't have been anyone who's in here. I know.\",\n",
       "  'src_lang': 'it',\n",
       "  'tgt_lang': 'en'},\n",
       " 'evaluation': {'Accuracy Issues': [],\n",
       "  'Readability Issues': [],\n",
       "  'Accuracy Score': '4',\n",
       "  'Readability Score': '4',\n",
       "  'Confidence Level': 'the_translation_is_excellent_without_any_error_spans_and_no_creative_liberties_were_taken',\n",
       "  'Main Vs Alternate': 'Alternate Translated Text has marginally better quality',\n",
       "  'Score': 32}}"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_path = \"/root/notebooks/MT_TQ/TQ/DataPrep_Prompting_Experiments/labeled_data/parsed/\"\n",
    "json_files = [os.path.join(root, file) for root, _, files in os.walk(data_path) for file in files if file.endswith('.json') and 'PLDL' in file]\n",
    "\n",
    "training_samples = []\n",
    "for json_file in tqdm(json_files):\n",
    "    with open(json_file, 'r') as file:\n",
    "        data = json.load(file)\n",
    "    sampled_items = random.sample(data[\"data\"], 20)\n",
    "    training_samples.extend(sampled_items)\n",
    "\n",
    "datapoints = []\n",
    "\n",
    "for sample in training_samples:\n",
    "    datapoint = {\"input\":{}}\n",
    "    datapoint[\"input\"][\"src_text\"] = sample[\"main_src_text\"]\n",
    "    datapoint[\"input\"][\"tgt_text\"] = sample[\"tgt_text\"]\n",
    "    datapoint[\"input\"][\"src_prev\"] = sample[\"tt_src_prev\"]\n",
    "    datapoint[\"input\"][\"src_next\"] = sample[\"tt_src_next\"]\n",
    "    datapoint[\"input\"][\"tgt_prev\"] = sample[\"tt_tgt_prev\"]\n",
    "    datapoint[\"input\"][\"tgt_next\"] = sample[\"tt_tgt_next\"]\n",
    "    datapoint[\"input\"][\"src_lang\"] = sample[\"src_lang\"]\n",
    "    datapoint[\"input\"][\"tgt_lang\"] = sample[\"tgt_lang\"]\n",
    "    datapoint[\"evaluation\"] = sample[\"labelers\"][0][\"annotation\"]\n",
    "    datapoints.append(datapoint)\n",
    "\n",
    "datapoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "bde34303-2f52-415f-b117-264e266b84f0",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-04-22T23:04:16.302953Z",
     "iopub.status.busy": "2025-04-22T23:04:16.302402Z",
     "iopub.status.idle": "2025-04-22T23:04:16.334330Z",
     "shell.execute_reply": "2025-04-22T23:04:16.333644Z",
     "shell.execute_reply.started": "2025-04-22T23:04:16.302928Z"
    }
   },
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "module 'dspy' has no attribute 'Predictor'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[35], line 28\u001b[0m\n\u001b[1;32m     25\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m prediction\u001b[38;5;241m.\u001b[39mevaluation\n\u001b[1;32m     27\u001b[0m \u001b[38;5;66;03m# Create a custom predictor using your Netflix model\u001b[39;00m\n\u001b[0;32m---> 28\u001b[0m \u001b[38;5;28;01mclass\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mNetflixPredictor\u001b[39;00m(\u001b[43mdspy\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mPredictor\u001b[49m):\n\u001b[1;32m     29\u001b[0m     \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21m__init__\u001b[39m(\u001b[38;5;28mself\u001b[39m, model):\n\u001b[1;32m     30\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmodel \u001b[38;5;241m=\u001b[39m model\n",
      "\u001b[0;31mAttributeError\u001b[0m: module 'dspy' has no attribute 'Predictor'"
     ]
    }
   ],
   "source": [
    "class TranslationQualityChecker(dspy.Signature):\n",
    "    \"\"\"Evaluate the quality of translation.\"\"\"\n",
    "    \n",
    "    context = dspy.InputField(desc=\"Source and target text with context\")\n",
    "    evaluation = dspy.OutputField(desc=\"Detailed evaluation of the translation quality\")\n",
    "\n",
    "class TranslationQualityModule(dspy.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.checker = dspy.Predict(TranslationQualityChecker)\n",
    "    \n",
    "    def forward(self, src_text, tgt_text, src_prev, tgt_prev, src_next, tgt_next, src_lang, tgt_lang):\n",
    "        context = {\n",
    "            \"source_text\": src_text,\n",
    "            \"target_text\": tgt_text,\n",
    "            \"source_previous\": src_prev,\n",
    "            \"target_previous\": tgt_prev,\n",
    "            \"source_next\": src_next,\n",
    "            \"target_next\": tgt_next,\n",
    "            \"source_language\": src_lang,\n",
    "            \"target_language\": tgt_lang\n",
    "        }\n",
    "        \n",
    "        prediction = self.checker(context=context)\n",
    "        return prediction.evaluation\n",
    "\n",
    "# Create a custom backend using your Netflix model\n",
    "class NetflixBackend(dspy.BackendBase):\n",
    "    def __init__(self, model):\n",
    "        super().__init__()\n",
    "        self.model = model\n",
    "        \n",
    "    def complete(self, prompt, **kwargs):\n",
    "        messages = [{\"role\": \"user\", \"content\": prompt}]\n",
    "        response = self.model.generate(messages)\n",
    "        return response\n",
    "\n",
    "    def completions(self, prompts, **kwargs):\n",
    "        return [self.complete(prompt, **kwargs) for prompt in prompts]\n",
    "\n",
    "# Prepare training data\n",
    "def prepare_training_data(data_points):\n",
    "    compiled_data = []\n",
    "    for dp in data_points:\n",
    "        input_data = dp['input']\n",
    "        train_example = dspy.Example(\n",
    "            context={\n",
    "                \"source_text\": input_data['src_text'],\n",
    "                \"target_text\": input_data['tgt_text'],\n",
    "                \"source_previous\": input_data['src_prev'],\n",
    "                \"target_previous\": input_data['tgt_prev'],\n",
    "                \"source_next\": input_data['src_next'],\n",
    "                \"target_next\": input_data['tgt_next'],\n",
    "                \"source_language\": input_data['src_lang'],\n",
    "                \"target_language\": input_data['tgt_lang']\n",
    "            },\n",
    "            evaluation=dp['evaluation']\n",
    "        )\n",
    "        compiled_data.append(train_example)\n",
    "    return compiled_data\n",
    "\n",
    "def optimize_prompt(model, training_data, validation_data):\n",
    "    # Initialize DSPy with your custom backend\n",
    "    backend = NetflixBackend(model)\n",
    "    dspy.settings.configure(lm=backend)\n",
    "    \n",
    "    # Create the optimizer\n",
    "    optimizer = MIPROv2(\n",
    "        metric=\"exact_match\",  # or another appropriate metric\n",
    "        max_rounds=5,\n",
    "        max_prompts=3,\n",
    "        temp=0.7\n",
    "    )\n",
    "    \n",
    "    # Compile the module\n",
    "    translation_module = TranslationQualityModule()\n",
    "    \n",
    "    # Optimize the prompt\n",
    "    optimized_module = optimizer.optimize(\n",
    "        module=translation_module,\n",
    "        trainset=training_data,\n",
    "        valset=validation_data,\n",
    "        metric=dspy.evaluate.answer_exact_match\n",
    "    )\n",
    "    \n",
    "    return optimized_module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67a4583f-162c-4e2d-b061-798f6c676a28",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TranslationQualityAssessor(dspy.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.assess = dspy.ChainOfThought(TranslationQualitySignature)\n",
    "\n",
    "    def forward(self, src_lang, tgt_lang, src_text, translation, src_prev=\"\", tgt_prev=\"\", src_next=\"\", tgt_next=\"\"):\n",
    "        context = f\"\"\"Previous Context:\n",
    "                Source: {src_prev}\n",
    "                Translation: {tgt_prev}\n",
    "                \n",
    "                Next Context:\n",
    "                Source: {src_next}\n",
    "                Translation: {tgt_next}\"\"\"\n",
    "\n",
    "        result = self.assess(\n",
    "            context=context,\n",
    "            source=f\"Source ({src_lang}): {src_text}\",\n",
    "            translation=f\"Translation ({tgt_lang}): {translation}\"\n",
    "        )\n",
    "        \n",
    "        return result.evaluation\n",
    "\n",
    "class TranslationMetrics:\n",
    "    @staticmethod\n",
    "    def exact_match_score(pred, gold):\n",
    "        try:\n",
    "            pred_json = json.loads(pred)\n",
    "            gold_json = gold\n",
    "            \n",
    "            accuracy_match = (str(pred_json.get('Accuracy Score')) == str(gold_json.get('Accuracy Score')))\n",
    "            readability_match = (str(pred_json.get('Readability Score')) == str(gold_json.get('Readability Score')))\n",
    "            \n",
    "            return (accuracy_match and readability_match)\n",
    "        except:\n",
    "            return False\n",
    "    \n",
    "    @staticmethod\n",
    "    def partial_match_score(pred, gold):\n",
    "        try:\n",
    "            pred_json = json.loads(pred)\n",
    "            gold_json = gold\n",
    "            \n",
    "            # Score comparison\n",
    "            accuracy_diff = abs(float(pred_json.get('Accuracy Score', 0)) - float(gold_json.get('Accuracy Score', 0)))\n",
    "            readability_diff = abs(float(pred_json.get('Readability Score', 0)) - float(gold_json.get('Readability Score', 0)))\n",
    "            \n",
    "            # Issues comparison\n",
    "            pred_accuracy_issues = set(str(issue) for issue in pred_json.get('Accuracy Issues', []))\n",
    "            gold_accuracy_issues = set(str(issue) for issue in gold_json.get('Accuracy Issues', []))\n",
    "            pred_readability_issues = set(str(issue) for issue in pred_json.get('Readability Issues', []))\n",
    "            gold_readability_issues = set(str(issue) for issue in gold_json.get('Readability Issues', []))\n",
    "            \n",
    "            # Calculate Jaccard similarity for issues\n",
    "            accuracy_issues_sim = len(pred_accuracy_issues & gold_accuracy_issues) / max(1, len(pred_accuracy_issues | gold_accuracy_issues))\n",
    "            readability_issues_sim = len(pred_readability_issues & gold_readability_issues) / max(1, len(pred_readability_issues | gold_readability_issues))\n",
    "            \n",
    "            # Combine scores (0.6 weight to scores, 0.4 to issues similarity)\n",
    "            score_component = 1 - ((accuracy_diff + readability_diff) / 8)\n",
    "            issues_component = (accuracy_issues_sim + readability_issues_sim) / 2\n",
    "            \n",
    "            final_score = 0.6 * score_component + 0.4 * issues_component\n",
    "            return max(0, final_score)\n",
    "        except:\n",
    "            return 0\n",
    "\n",
    "def prepare_dataset(file_path):\n",
    "    with open(file_path, 'r') as f:\n",
    "        data = json.load(f)\n",
    "    \n",
    "    prepared_data = []\n",
    "    \n",
    "    for item in data:\n",
    "        example = dspy.Example(\n",
    "            context=f\"\"\"Previous Context:\n",
    "                    Source: {item['src_prev']}\n",
    "                    Translation: {item['tgt_prev']}\n",
    "                    \n",
    "                    Next Context:\n",
    "                    Source: {item['src_next']}\n",
    "                    Translation: {item['tgt_next']}\"\"\",\n",
    "            source=f\"Source ({item['src_lang']}): {item['src_text']}\",\n",
    "            translation=f\"Translation ({item['tgt_lang']}): {item['main_text']}\",\n",
    "            evaluation=json.dumps(item['evaluation'], ensure_ascii=False)\n",
    "        ).with_inputs(\"context\", \"source\", \"translation\")\n",
    "        \n",
    "        prepared_data.append(example)\n",
    "    \n",
    "    # Split data: 70% train, 15% dev, 15% test\n",
    "    train_size = int(0.7 * len(prepared_data))\n",
    "    dev_size = int(0.15 * len(prepared_data))\n",
    "    \n",
    "    train_data = prepared_data[:train_size]\n",
    "    dev_data = prepared_data[train_size:train_size + dev_size]\n",
    "    test_data = prepared_data[train_size + dev_size:]\n",
    "    \n",
    "    return train_data, dev_data, test_data\n",
    "\n",
    "def optimize_translation_quality_assessment():\n",
    "    # Initialize DSPy\n",
    "    lm = TranslationQualityLM()\n",
    "    dspy.settings.configure(lm=lm)\n",
    "    \n",
    "    # Load and prepare dataset\n",
    "    train_data, dev_data, test_data = prepare_dataset('translation_quality_dataset.json')\n",
    "    \n",
    "    # Create evaluator\n",
    "    evaluator = Evaluate(\n",
    "        metrics={\n",
    "            'exact_match': TranslationMetrics.exact_match_score,\n",
    "            'partial_match': TranslationMetrics.partial_match_score\n",
    "        }\n",
    "    )\n",
    "    \n",
    "    # Initialize module\n",
    "    assessor = TranslationQualityAssessor()\n",
    "    \n",
    "    # Initialize MIPROv2 optimizer\n",
    "    optimizer = dspy.MIPROv2(\n",
    "        metric=lambda x: x['partial_match'],\n",
    "        max_rounds=5,              # Number of optimization rounds\n",
    "        max_traces=10,            # Number of traces per round\n",
    "        max_depth=3,              # Maximum depth of reasoning chains\n",
    "        num_candidate_prompts=5,  # Number of candidate prompts to generate\n",
    "        num_rounds_per_prompt=3,  # Number of rounds per candidate prompt\n",
    "        temperature=0.7,\n",
    "        verbose=True\n",
    "    )\n",
    "    \n",
    "    # Compile the module with optimization\n",
    "    compiled_assessor = optimizer.compile(\n",
    "        assessor,\n",
    "        trainset=train_data,\n",
    "        devset=dev_data,\n",
    "        eval_kwargs={\n",
    "            'metric': 'partial_match',\n",
    "            'num_threads': 4,\n",
    "            'batch_size': 8\n",
    "        }\n",
    "    )\n",
    "    \n",
    "    # Evaluate on test set\n",
    "    results = []\n",
    "    for example in test_data:\n",
    "        pred = compiled_assessor(\n",
    "            context=example.context,\n",
    "            source=example.source,\n",
    "            translation=example.translation\n",
    "        )\n",
    "        \n",
    "        result = evaluator.evaluate(\n",
    "            predictions=[pred],\n",
    "            ground_truth=[example.evaluation]\n",
    "        )\n",
    "        results.append(result)\n",
    "    \n",
    "    # Calculate and print final metrics\n",
    "    avg_exact_match = np.mean([r['exact_match'] for r in results])\n",
    "    avg_partial_match = np.mean([r['partial_match'] for r in results])\n",
    "    \n",
    "    print(f\"Average Exact Match Score: {avg_exact_match:.3f}\")\n",
    "    print(f\"Average Partial Match Score: {avg_partial_match:.3f}\")\n",
    "    \n",
    "    return compiled_assessor\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    optimized_assessor = optimize_translation_quality_assessment()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "timedlibs",
   "language": "python",
   "name": "timedlibs"
  },
  "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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}