{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "55207dc8-0fcc-4ada-8b3e-80bb98df82c1",
   "metadata": {},
   "source": [
    "# Gemini Usage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "ad6bf125-d5df-4d01-ad67-37db5b70e85b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The history saving thread hit an unexpected error (OperationalError('attempt to write a readonly database')).History will not be written to the database.\n"
     ]
    }
   ],
   "source": [
    "# !pip install google-generativeai==0.8.3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a5c56369-5d21-4b0d-aa74-8b09d190e857",
   "metadata": {},
   "outputs": [],
   "source": [
    "import google.generativeai as genai\n",
    "import json\n",
    "from utils import load_prompts\n",
    "prompts = load_prompts(\"prompts.yaml\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d203c836-5e1e-4145-a282-686e4aee5f68",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'system': {'general': 'You are an expert at Natural Language Inference (NLI). Your task is to analyze two pieces of text - a premise and a hypothesis - and classify them into NEUTRAL, ENTAILMENT, or CONTRADICTION.'}, 'user': {'general': 'premise: {PREMISE_HERE}, hypothesis: {HYPOTHESIS_HERE}.', 'json_mode': \"Generate the classification result in JSON with the key: 'result': str (NEUTRAL, ENTAILMENT, or CONTRADICTION)\", 'few': 'USER: premise: {PREMISE_1}, hypothesis: {HYPOTHESIS_1}.\\nMODEL: {RESULT_1}\\nUSER: premise: {PREMISE_2}, hypothesis: {HYPOTHESIS_2}.\\nMODEL: {RESULT_2}\\nUSER: premise: {PREMISE_3}, hypothesis: {HYPOTHESIS_3}.\\nMODEL:'}, 'mutual': {'few_hint': 'Please first check the following examples.'}}\n"
     ]
    }
   ],
   "source": [
    "print(prompts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c10e382f-a1ca-47db-9711-5ae8ce78ef7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "MY_GOOGLE_API_KEY = \"MY_GOOGLE_API_KEY\"\n",
    "genai.configure(api_key=MY_GOOGLE_API_KEY)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "34334b7d-32f2-4cc5-bb6d-f9d240d2d695",
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_NAME = \"gemini-1.5-pro\"\n",
    "TEMPERATURE = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be4fb61a-6897-4afa-aada-13563e1decaa",
   "metadata": {},
   "source": [
    "## Example for Classification\n",
    "- Dataset: SemEval 2014 Task1-2 (3-class classification)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0dae8347-bcc7-48d3-8cbf-a5f74d3bc7c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = {\n",
    "    \"premise\": \"A group of kids is playing in a yard and an old man is standing in the background\",\n",
    "    \"hypothesis\": \"A group of boys in a yard is playing and a man is standing in the background\", \n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8bcbf02e-6581-48e6-a38f-c51718660b42",
   "metadata": {},
   "outputs": [],
   "source": [
    "system_prompt = prompts[\"system\"][\"general\"]\n",
    "user_prompt = prompts[\"user\"][\"general\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "54ae7de1-63cb-45ee-a2bb-d4b44a407411",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "You are an expert at Natural Language Inference (NLI). Your task is to analyze two pieces of text - a premise and a hypothesis - and classify them into NEUTRAL, ENTAILMENT, or CONTRADICTION.\n",
      "premise: {PREMISE_HERE}, hypothesis: {HYPOTHESIS_HERE}.\n"
     ]
    }
   ],
   "source": [
    "print(system_prompt)\n",
    "print(user_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "5a48c773-8d2d-438b-88c0-2ae5d07d1f56",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = genai.GenerativeModel(\n",
    "    MODEL_NAME,\n",
    "    generation_config={\"temperature\": TEMPERATURE},\n",
    "    system_instruction=system_prompt,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "b2bec27e-f1a7-406f-9232-ac3d494f3de6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "premise: A group of kids is playing in a yard and an old man is standing in the background, hypothesis: A group of boys in a yard is playing and a man is standing in the background.\n"
     ]
    }
   ],
   "source": [
    "cur_user_prompt = user_prompt.format(\n",
    "    PREMISE_HERE=inputs[\"premise\"],\n",
    "    HYPOTHESIS_HERE=inputs[\"hypothesis\"]\n",
    ")\n",
    "print(cur_user_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d2a39d10-1317-47a4-b841-0e11f413f817",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = model.generate_content(cur_user_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "2e9c211a-e7eb-40b1-88d7-26ee59951f51",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NEUTRAL.  The premise states \"kids,\" which could include girls. The hypothesis specifies \"boys.\"  The premise says \"old man,\" while the hypothesis just says \"man.\"  While an old man *is* a man, the hypothesis doesn't require the man to be old.  Therefore, the hypothesis is neither entailed nor contradicted by the premise.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(response.text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "46110d8f-a78e-47d8-9e68-884bb18f9f71",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Num of input tokens: 173\n"
     ]
    }
   ],
   "source": [
    "num_sys_tokens = str(model.count_tokens(system_prompt)).split(\": \")[1].strip()\n",
    "num_usr_tokens = str(model.count_tokens(cur_user_prompt)).split(\": \")[1].strip()\n",
    "num_input_tokens = int(num_sys_tokens) + int(num_usr_tokens)\n",
    "print(f\"Num of input tokens: {num_input_tokens}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "5abb7ead-f79b-46d5-b98b-eaace44a7c59",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Num of output tokens 119\n"
     ]
    }
   ],
   "source": [
    "num_output_tokens = str(model.count_tokens(response.text)).split(\": \")[1].strip()\n",
    "print(f\"Num of output tokens {num_output_tokens}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3aae09fa-ed4c-4c13-946c-3222f94dd014",
   "metadata": {},
   "source": [
    "### Generate structured output in JSON"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "cc69ea10-19e1-4633-90f3-dedcdc46fd62",
   "metadata": {},
   "outputs": [],
   "source": [
    "user_prompt_json = prompts[\"user\"][\"general\"] + \" \" + prompts[\"user\"][\"json_mode\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "a5b741b2-5500-46d7-a480-db14432b6678",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "premise: A group of kids is playing in a yard and an old man is standing in the background, hypothesis: A group of boys in a yard is playing and a man is standing in the background. Generate the classification result in JSON with the key: 'result': str (NEUTRAL, ENTAILMENT, or CONTRADICTION)\n"
     ]
    }
   ],
   "source": [
    "cur_user_prompt = user_prompt_json.format(\n",
    "    PREMISE_HERE=inputs[\"premise\"],\n",
    "    HYPOTHESIS_HERE=inputs[\"hypothesis\"]\n",
    ")\n",
    "print(cur_user_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "3307a22c-2413-400a-ab61-c9af4cd6fb82",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = model.generate_content(\n",
    "    cur_user_prompt,\n",
    "    generation_config={\n",
    "        \"response_mime_type\": \"application/json\",\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "72f1bd25-faf5-4087-a37a-5cb595c2e591",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "str"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(response.text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "d57965af-fd26-4f4f-a5b4-68c3e3d1d7ca",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'dict'>\n",
      "{'result': 'NEUTRAL'}\n"
     ]
    }
   ],
   "source": [
    "result_json = json.loads(response.text)\n",
    "print(type(result_json))\n",
    "print(result_json)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba5bfae9-723a-48a4-b52c-f3c740b55438",
   "metadata": {},
   "source": [
    "### Few-shot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "f705b761-cc2e-4e23-a956-644341f1ad1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "fs_user_prompt = prompts[\"user\"][\"few\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "6275f3b6-b39a-4920-9823-c1ea70e12253",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "USER: premise: A group of kids is playing in a yard and an old man is standing in the background, hypothesis: A group of boys in a yard is playing and a man is standing in the background.\n",
      "MODEL: {'result': 'NEUTRAL'}\n",
      "USER: premise: A man, a woman and two girls are walking on the beach, hypothesis: A group of people is on a beach.\n",
      "MODEL: {'result': 'ENTAILMENT'}\n",
      "USER: premise: Two teams are competing in a football match, hypothesis: Two groups of people are playing football.\n",
      "MODEL:\n"
     ]
    }
   ],
   "source": [
    "cur_fs_user_prompt = fs_user_prompt.format(\n",
    "    PREMISE_1=\"A group of kids is playing in a yard and an old man is standing in the background\",\n",
    "    HYPOTHESIS_1=\"A group of boys in a yard is playing and a man is standing in the background\",\n",
    "    RESULT_1=\"{'result': 'NEUTRAL'}\",\n",
    "    PREMISE_2=\"A man, a woman and two girls are walking on the beach\",\n",
    "    HYPOTHESIS_2=\"A group of people is on a beach\",\n",
    "    RESULT_2=\"{'result': 'ENTAILMENT'}\",\n",
    "    PREMISE_3=\"Two teams are competing in a football match\",\n",
    "    HYPOTHESIS_3=\"Two groups of people are playing football\",  \n",
    ")\n",
    "print(cur_fs_user_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "cb495bf0-5ee1-4e6e-b1e0-f050b08a7062",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "You are an expert at Natural Language Inference (NLI). Your task is to analyze two pieces of text - a premise and a hypothesis - and classify them into NEUTRAL, ENTAILMENT, or CONTRADICTION. Please first check the following examples.\n"
     ]
    }
   ],
   "source": [
    "fs_system_prompt = prompts[\"system\"][\"general\"] + \" \" + prompts[\"mutual\"][\"few_hint\"]\n",
    "print(fs_system_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "bf10cd01-2ceb-4a4a-b22e-f2ce6defbbe0",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = genai.GenerativeModel(\n",
    "    MODEL_NAME,\n",
    "    generation_config={\"temperature\": TEMPERATURE},\n",
    "    system_instruction=fs_system_prompt,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "e4e2783f-1b02-4209-abfc-f5b74713f1e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = model.generate_content(\n",
    "    cur_fs_user_prompt,\n",
    "    generation_config={\n",
    "        \"response_mime_type\": \"application/json\",\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "1034be34-c78d-4681-a8bd-72106ff214c1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'result': 'ENTAILMENT'}\n"
     ]
    }
   ],
   "source": [
    "result_json = json.loads(response.text)\n",
    "print(result_json)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e00c92c-31fb-4a4d-8a01-2ee20f86164e",
   "metadata": {},
   "source": [
    "## Example for Summarization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "59b1573a-6ec3-4c67-bd6f-c19cbd1fe78b",
   "metadata": {},
   "outputs": [],
   "source": [
    "system_prompt = \"\"\"你是個中文文本摘要的專家，現在請你對一篇輸入的文章進行摘要。\"\"\"\n",
    "\n",
    "model = genai.GenerativeModel(\n",
    "    MODEL_NAME,\n",
    "    generation_config={\"temperature\": TEMPERATURE},\n",
    "    system_instruction=system_prompt,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "6f991b65-b6ce-4b1a-afa2-df876142e4a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_source_txt = \"新华社受权于18日全文播发修改后的《中华人民共和国立法法》，修改后的立法法分为“总则”“法律”“行政法规”“地方性法规、自治条例和单行条例、规章”“适用与备案审查”“附则”等6章，共计105条。\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "cbdbbff9-f926-439b-a914-a293bc67c0d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = model.generate_content(input_source_txt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "947de4bf-e71f-400e-bfdd-c5e17b6ce125",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'新修订的《中华人民共和国立法法》已正式发布，共六章105条，涵盖总则、法律、行政法规、地方性法规及规章、适用与备案审查以及附则等方面。\\n'"
      ]
     },
     "execution_count": 66,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(response.text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2f6f7c1-6b71-437f-8e72-a72ea0f91192",
   "metadata": {},
   "source": [
    "# Claude Usage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e202ca1-5a65-48de-8631-1afebfdfaaa3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install anthropic==0.39.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0a7b1a0f-9bf3-4967-af30-97d2f3b468f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import anthropic\n",
    "import json\n",
    "import re\n",
    "from utils import load_prompts\n",
    "prompts = load_prompts(\"prompts.yaml\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e7c9aba8-bf01-40fa-9579-cfff450eefba",
   "metadata": {},
   "outputs": [],
   "source": [
    "MY_anthropic_KEY = \"MY_anthropic_KEY\"\n",
    "client = anthropic.Anthropic(api_key=MY_anthropic_KEY)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b02cec37-1cc9-46ce-82de-387519362abc",
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_NAME = \"claude-3-5-sonnet-20241022\"\n",
    "TEMPERATURE = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "feb1b843-a039-4946-a56a-a65f5cff598e",
   "metadata": {},
   "outputs": [],
   "source": [
    "system_prompt = prompts[\"system\"][\"general\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "17f62cc8-1968-4033-a439-4d8abb83cea8",
   "metadata": {},
   "outputs": [],
   "source": [
    "user_prompt_json = prompts[\"user\"][\"general\"] + \" \" + prompts[\"user\"][\"json_mode\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "07163c7b-84be-41db-93e0-fe30b4d91f07",
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = {\n",
    "    \"premise\": \"A group of kids is playing in a yard and an old man is standing in the background\",\n",
    "    \"hypothesis\": \"A group of boys in a yard is playing and a man is standing in the background\", \n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "45f8a768-2568-4fff-81e5-4d2f1d9cca4f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "premise: A group of kids is playing in a yard and an old man is standing in the background, hypothesis: A group of boys in a yard is playing and a man is standing in the background. Generate the classification result in JSON with the key: 'result': str (NEUTRAL, ENTAILMENT, or CONTRADICTION)\n"
     ]
    }
   ],
   "source": [
    "cur_user_prompt = user_prompt_json.format(\n",
    "    PREMISE_HERE=inputs[\"premise\"],\n",
    "    HYPOTHESIS_HERE=inputs[\"hypothesis\"]\n",
    ")\n",
    "print(cur_user_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "bd6c8d8d-06c1-4a7f-806f-673c25830cdc",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = client.messages.create(\n",
    "    model=MODEL_NAME,\n",
    "    max_tokens=1000,\n",
    "    temperature=TEMPERATURE,\n",
    "    system=system_prompt,\n",
    "    messages=[\n",
    "        {\n",
    "            \"role\": \"user\",\n",
    "            \"content\": [{\"type\": \"text\", \"text\": cur_user_prompt}],\n",
    "        },\n",
    "    ],\n",
    ")\n",
    "result = response.content[0].text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "dcfbdb5a-06d0-4f92-ba18-8ef5c19f8975",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Let me analyze this carefully:\\n\\n1. Premise mentions \"kids\" (gender-neutral) while hypothesis specifies \"boys\" (male)\\n2. Premise mentions \"old man\" while hypothesis just says \"man\"\\n3. Both texts agree on:\\n   - Group playing in a yard\\n   - Man standing in background\\n\\nSince the hypothesis specifies \"boys\" when the premise only mentions \"kids\" (which could be boys, girls, or mixed), and drops the \"old\" qualifier for the man, this is a NEUTRAL case. The hypothesis could be true, but it\\'s not necessarily entailed by the premise.\\n\\n{\\n    \"result\": \"NEUTRAL\"\\n}'"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "3c034853-add6-424e-9b98-8825a0b8ea7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract the JSON-like part of the response using regex to match a JSON object\n",
    "json_match = re.search(r\"\\{.*?\\}\", response.content[0].text, re.DOTALL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "1bf9db36-a4e7-48b2-8c4a-863a27af5616",
   "metadata": {},
   "outputs": [],
   "source": [
    "json_str = json_match.group(0)  # Extract the JSON string\n",
    "# Now parse the extracted JSON\n",
    "result = json.loads(json_str, strict=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "231b74a4-1e35-4628-8939-096bf53b22b9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'result': 'NEUTRAL'}"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "333a982b-54eb-468a-a5ba-05c30a9e41f2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Num of input tokens 122\n",
      "Num of output tokens 149\n"
     ]
    }
   ],
   "source": [
    "print(f\"Num of input tokens {response.usage.input_tokens}\")\n",
    "print(f\"Num of output tokens {response.usage.output_tokens}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "580640e4-8a03-48a8-a66f-f9dae9e53939",
   "metadata": {},
   "source": [
    "## Few-shot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "21a84c8e-75f7-4df0-9ee5-9e2f66719e96",
   "metadata": {},
   "outputs": [],
   "source": [
    "fs_system_prompt = prompts[\"system\"][\"general\"] + \" \" + prompts[\"mutual\"][\"few_hint\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "f46281fc-54be-403e-9cde-30e257631c6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "fs_user_prompt = prompts[\"user\"][\"few\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "6ce2fb20-c888-4dc9-8082-555ee922ec47",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "USER: premise: A group of kids is playing in a yard and an old man is standing in the background, hypothesis: A group of boys in a yard is playing and a man is standing in the background.\n",
      "MODEL: {'result': 'NEUTRAL'}\n",
      "USER: premise: A man, a woman and two girls are walking on the beach, hypothesis: A group of people is on a beach.\n",
      "MODEL: {'result': 'ENTAILMENT'}\n",
      "USER: premise: Two teams are competing in a football match, hypothesis: Two groups of people are playing football.\n",
      "MODEL:\n"
     ]
    }
   ],
   "source": [
    "cur_fs_user_prompt = fs_user_prompt.format(\n",
    "    PREMISE_1=\"A group of kids is playing in a yard and an old man is standing in the background\",\n",
    "    HYPOTHESIS_1=\"A group of boys in a yard is playing and a man is standing in the background\",\n",
    "    RESULT_1=\"{'result': 'NEUTRAL'}\",\n",
    "    PREMISE_2=\"A man, a woman and two girls are walking on the beach\",\n",
    "    HYPOTHESIS_2=\"A group of people is on a beach\",\n",
    "    RESULT_2=\"{'result': 'ENTAILMENT'}\",\n",
    "    PREMISE_3=\"Two teams are competing in a football match\",\n",
    "    HYPOTHESIS_3=\"Two groups of people are playing football\",  \n",
    ")\n",
    "print(cur_fs_user_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "efd7064e-0b1d-4f10-b25a-1d6ccffacf2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = client.messages.create(\n",
    "    model=MODEL_NAME,\n",
    "    max_tokens=1000,\n",
    "    temperature=TEMPERATURE,\n",
    "    system=fs_system_prompt,\n",
    "    messages=[\n",
    "        {\n",
    "            \"role\": \"user\",\n",
    "            \"content\": [{\"type\": \"text\", \"text\": cur_user_prompt}],\n",
    "        },\n",
    "    ],\n",
    ")\n",
    "result = response.content[0].text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "bbc5a95c-89a8-40a2-ad8e-16b0e5cddeb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract the JSON-like part of the response using regex to match a JSON object\n",
    "json_match = re.search(r\"\\{.*?\\}\", response.content[0].text, re.DOTALL)\n",
    "json_str = json_match.group(0)  # Extract the JSON string\n",
    "# Now parse the extracted JSON\n",
    "result = json.loads(json_str, strict=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "a989aaab-bb13-4a37-8c05-cb22fd012699",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'result': 'NEUTRAL'}"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e32f1f50-ae8f-477e-88bc-c0ccc678aaa2",
   "metadata": {},
   "source": [
    "# OpenAI Usage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75d968cf-2066-46d1-b6fe-834ea2d9eedf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install openai==1.54.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5faf8051-0fc4-4dae-89bd-5682fc6d966f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from openai import OpenAI\n",
    "import json\n",
    "from utils import load_prompts\n",
    "prompts = load_prompts(\"prompts.yaml\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e607cf57-ae84-4360-bfc5-d0bfeb704189",
   "metadata": {},
   "outputs": [],
   "source": [
    "MY_openai_KEY = \"MY_openai_KEY\"\n",
    "client = OpenAI(api_key=MY_openai_KEY)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "021eccf5-d9bd-4cb8-94e6-19875dfb65a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_NAME = \"gpt-4o\"\n",
    "TEMPERATURE = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3b9c5331-b062-4b9c-b52e-e1ffd87b2aea",
   "metadata": {},
   "outputs": [],
   "source": [
    "system_prompt = prompts[\"system\"][\"general\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "339718ec-0b94-440b-b394-1195890aebe4",
   "metadata": {},
   "outputs": [],
   "source": [
    "user_prompt_json = prompts[\"user\"][\"general\"] + \" \" + prompts[\"user\"][\"json_mode\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f78cca64-ef10-4de8-ae06-fb5d31ac95da",
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = {\n",
    "    \"premise\": \"A group of kids is playing in a yard and an old man is standing in the background\",\n",
    "    \"hypothesis\": \"A group of boys in a yard is playing and a man is standing in the background\", \n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "39021c21-8b43-45b2-8501-0de210687b7c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "premise: A group of kids is playing in a yard and an old man is standing in the background, hypothesis: A group of boys in a yard is playing and a man is standing in the background. Generate the classification result in JSON with the key: 'result': str (NEUTRAL, ENTAILMENT, or CONTRADICTION)\n"
     ]
    }
   ],
   "source": [
    "cur_user_prompt = user_prompt_json.format(\n",
    "    PREMISE_HERE=inputs[\"premise\"],\n",
    "    HYPOTHESIS_HERE=inputs[\"hypothesis\"]\n",
    ")\n",
    "print(cur_user_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3da88401-fb14-45c2-81c5-6164cf9fbab1",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = client.chat.completions.create(\n",
    "    model=MODEL_NAME,\n",
    "    response_format={\"type\": \"json_object\"},\n",
    "    messages=[\n",
    "        {\"role\": \"system\", \"content\": system_prompt},\n",
    "        {\"role\": \"user\", \"content\": cur_user_prompt},\n",
    "    ],\n",
    "    temperature=TEMPERATURE,\n",
    "    # max_tokens=max_tokens,\n",
    ")\n",
    "result_json = json.loads(response.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "cf8de3b4-d2ea-436d-9cec-0642447f1ba6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'result': 'ENTAILMENT'}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result_json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e46b154-862e-4e95-a92d-c331deebe3dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Num of input tokens {response.usage.prompt_tokens}\")\n",
    "print(f\"Num of output tokens {response.usage.completion_tokens}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b32e982a-21d7-4148-85d8-e0144f223830",
   "metadata": {},
   "source": [
    "### Few-shot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "066491c8-2ab6-477c-907b-10191d915055",
   "metadata": {},
   "outputs": [],
   "source": [
    "fs_system_prompt = prompts[\"system\"][\"general\"] + \" \" + prompts[\"mutual\"][\"few_hint\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "3b3b910f-055b-46f1-a586-53e9628f6fd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# For OpenAI API, when using response_format={\"type\": \"json_object\"},\n",
    "# we must tell the model to output in JSON in the user prompts.\n",
    "user_prompt = prompts[\"user\"][\"general\"] + \" \" + prompts[\"user\"][\"json_mode\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "231c1b6c-ad68-4fb5-8986-0692508c0905",
   "metadata": {},
   "outputs": [],
   "source": [
    "few_shot_inputs = [\n",
    "    [\n",
    "        \"A group of kids is playing in a yard and an old man is standing in the background\",\n",
    "        \"A group of boys in a yard is playing and a man is standing in the background\"\n",
    "    ],\n",
    "    [\n",
    "        \"A man, a woman and two girls are walking on the beach\",\n",
    "        \"A group of people is on a beach\"\n",
    "    ],\n",
    "]\n",
    "few_shot_labels = [\"NEUTRAL\", \"ENTAILMENT\"]\n",
    "\n",
    "# inputs without label\n",
    "real_inputs = [\n",
    "        \"Two teams are competing in a football match\",\n",
    "        \"Two groups of people are playing football\"\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "9fd980a3-39e7-48a7-8204-e489192d7a9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "messages = [{\"role\": \"system\", \"content\": fs_system_prompt}]\n",
    "for sents, label in zip(few_shot_inputs, few_shot_labels):\n",
    "    cur_user_prompt = user_prompt.format(\n",
    "        PREMISE_HERE=sents[0],\n",
    "        HYPOTHESIS_HERE=sents[1]\n",
    "    )\n",
    "    messages.append({\"role\": \"user\", \"content\": cur_user_prompt})\n",
    "    messages.append({\"role\": \"assistant\", \"content\": \"{\"+\"'result': \"+f\"{label}\"+\"}\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "12eee8d7-3c34-463c-b268-da2364c2809c",
   "metadata": {},
   "outputs": [],
   "source": [
    "cur_user_prompt = user_prompt.format(\n",
    "    PREMISE_HERE=real_inputs[0],\n",
    "    HYPOTHESIS_HERE=real_inputs[1]\n",
    ")\n",
    "messages.append({\"role\": \"user\", \"content\": cur_user_prompt})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "119fb4fb-504b-4c3d-89af-9afd7fe80e31",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'role': 'system',\n",
       "  'content': 'You are an expert at Natural Language Inference (NLI). Your task is to analyze two pieces of text - a premise and a hypothesis - and classify them into NEUTRAL, ENTAILMENT, or CONTRADICTION. Please first check the following examples.'},\n",
       " {'role': 'user',\n",
       "  'content': \"premise: A group of kids is playing in a yard and an old man is standing in the background, hypothesis: A group of boys in a yard is playing and a man is standing in the background. Generate the classification result in JSON with the key: 'result': str (NEUTRAL, ENTAILMENT, or CONTRADICTION)\"},\n",
       " {'role': 'assistant', 'content': \"{'result': NEUTRAL}\"},\n",
       " {'role': 'user',\n",
       "  'content': \"premise: A man, a woman and two girls are walking on the beach, hypothesis: A group of people is on a beach. Generate the classification result in JSON with the key: 'result': str (NEUTRAL, ENTAILMENT, or CONTRADICTION)\"},\n",
       " {'role': 'assistant', 'content': \"{'result': ENTAILMENT}\"},\n",
       " {'role': 'user',\n",
       "  'content': \"premise: Two teams are competing in a football match, hypothesis: Two groups of people are playing football. Generate the classification result in JSON with the key: 'result': str (NEUTRAL, ENTAILMENT, or CONTRADICTION)\"}]"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "8623f4f2-90f8-429c-8ee9-4686e165aa0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = client.chat.completions.create(\n",
    "    model=MODEL_NAME,\n",
    "    response_format={\"type\": \"json_object\"},\n",
    "    messages=messages,\n",
    "    temperature=TEMPERATURE,\n",
    "    # max_tokens=max_tokens,\n",
    ")\n",
    "result_json = json.loads(response.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "e1a78124-73dc-4897-97f5-72b996c70918",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'result': 'ENTAILMENT'}"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result_json"
   ]
  }
 ],
 "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.9.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}