{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "33c29df4-b030-403d-bc6e-a40911bddb15",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/dbamman/anlp25/blob/main/10.llms/HW9_LLM_Inference.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9efea054-c52a-4c71-8693-46ad38d235b0",
   "metadata": {},
   "source": [
    "# HW9: LLM Inference\n",
    "\n",
    "In this homework, you will experiment with different ways of improving LLM classification performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cf69bb7-561b-4946-a446-86d6e7381379",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "from tqdm import tqdm\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53f94f6d-16a0-49fe-9c5a-755c22fccb66",
   "metadata": {},
   "outputs": [],
   "source": [
    "# use the 4B model\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\"Qwen/Qwen3-4B\", device_map=\"cuda\", dtype=\"auto\")\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"Qwen/Qwen3-4B\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "678131e0-6c58-4372-8d74-71b7f3ab80ec",
   "metadata": {},
   "source": [
    "## Loading data\n",
    "\n",
    "For the majority of this homework, we will be using data from *Who Feels What and Why? Annotation of a Literature Corpus with Semantic Roles of Emotions* [(Kim and Klinger, 2018)](https://aclanthology.org/C18-1114.pdf)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d51a2d95-3c27-49c0-961a-7874b7113ae9",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://raw.githubusercontent.com/bamman-group/ca-classification-data/refs/heads/main/data/emotion/train.jsonl\n",
    "!wget https://raw.githubusercontent.com/bamman-group/ca-classification-data/refs/heads/main/data/emotion/test.jsonl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c8f4d02-bd14-4eef-9424-f9edfc13e65c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "def load_data(filepath):\n",
    "    with open(filepath, \"r\") as f:\n",
    "        data = [\n",
    "            json.loads(line) for line in f\n",
    "        ]\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc8ebcfe-5d8a-4027-bba8-2cf75afd0a32",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = load_data(\"train.jsonl\")\n",
    "test_data = load_data(\"test.jsonl\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5af8b60-c950-4282-b7b2-ec40fb428d65",
   "metadata": {},
   "source": [
    "### Question 1\n",
    "\n",
    "Take a look through the paper, as well as the actual dataset. What are the classification labels? **Fill them in below.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "456b1b93-31b6-4b84-95e7-07bea7563ee9",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36beeb2a-708b-405c-8d48-2f510ba8b496",
   "metadata": {},
   "outputs": [],
   "source": [
    "# FILL ME IN\n",
    "labels = [\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f00dff07-9811-43da-a7e4-25594978fc23",
   "metadata": {},
   "source": [
    "## Setting up the LLM\n",
    "\n",
    "For greater consistency, we set the temperature to a low value (0.01) by default, but this can be changed with the generation_config setting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cb2c455-f330-49fc-bc7f-07b25fa7c19c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from textwrap import dedent\n",
    "import itertools\n",
    "import inspect\n",
    "\n",
    "def call_llm(prompt, system_prompt=\"You are a helpful assistant.\", generation_config=None):  \n",
    "    if generation_config is None:\n",
    "        generation_config = {\n",
    "            \"max_new_tokens\": 10,\n",
    "            \"temperature\": 0.01\n",
    "        }\n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": system_prompt},\n",
    "        {\"role\": \"user\", \"content\": prompt}\n",
    "    ]\n",
    "    text = tokenizer.apply_chat_template(\n",
    "        messages,\n",
    "        tokenize=False,\n",
    "        add_generation_prompt=True,\n",
    "        enable_thinking=False\n",
    "    )\n",
    "\n",
    "    model_inputs = tokenizer([text], return_tensors=\"pt\").to(model.device)\n",
    "    \n",
    "    # conduct text completion\n",
    "    generated = model.generate(\n",
    "        **model_inputs,\n",
    "        **generation_config\n",
    "    )\n",
    "\n",
    "    # let's break this down:\n",
    "    #                      | we take the element of the batch (our batch size is 1)\n",
    "    #                      |  |-----------------------------| skip our original input\n",
    "    output_ids = generated[0][len(model_inputs.input_ids[0]):].tolist()\n",
    "\n",
    "    # decode into token space\n",
    "    return tokenizer.decode(output_ids, skip_special_tokens=True).strip(\"\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d3a0585-e7cf-4158-a623-9d2b31baf5c1",
   "metadata": {},
   "source": [
    "## Classification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75213031-efe7-4ee0-b26b-0d23b4e5412a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(classifier):\n",
    "    predictions = classifier(train_data, test_data)\n",
    "    return sum(pred == target[\"label\"] for pred, target in zip(predictions, test_data)) / len(test_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d09b2bae-cb57-4476-832f-475c07e2b7db",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import Counter\n",
    "from textwrap import dedent\n",
    "import random\n",
    "\n",
    "\n",
    "def classify_majority_label(train_data, test_data):\n",
    "    \"\"\"Majority label baseline\"\"\"\n",
    "\n",
    "    majority_class = Counter([d[\"label\"] for d in train_data]).most_common(1)[0][0]\n",
    "    for i, datum in enumerate(tqdm(test_data)):\n",
    "        test_predictions.append(majority_class)\n",
    "\n",
    "    return test_predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "014feaf1-4b75-46d0-b5cb-bcc0fbac1288",
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluate(classify_majority_label)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9666804f-00d0-425d-9028-d5ef12bf120e",
   "metadata": {},
   "source": [
    "Fill the rest of these in!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0277508e-59dc-4700-9b25-0bcbafafb30b",
   "metadata": {},
   "source": [
    "## Question 2\n",
    "\n",
    "We've implemented a majority vote baseline for you. \n",
    "\n",
    "Implement a zero-shot prompting classifier. Try at least 3 versions of the prompt to compare their outputs.\n",
    "\n",
    "**In a few sentences,** describe how different prompting choices result in different outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4db7ea91-a6d1-443b-9556-8daa6acffea9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def classify_zero_shot(train_data, test_data):\n",
    "    \"\"\"Classification with zero-shot prompting.\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6896962-0902-4736-bdf4-4eec5ae25891",
   "metadata": {},
   "source": [
    "## Question 3\n",
    "\n",
    "Implement the following:\n",
    "\n",
    "1. Few-shot (k=3) classification\n",
    "3. Zero-shot with chain-of-thought\n",
    "4. Few-shot (k=3) with chain of thought (you will need to write reasoning chains)\n",
    "5. Zero-shot with self-consistency (use `generation_config` to change the temperature)\n",
    "\n",
    "For each of these, print out the raw LLM output for the first 5 data points in the test data.\n",
    "\n",
    "Use the `evaluate` function to measure the accuracy of your method. **Write a few sentences comparing the performance of different prompting methods (including the above, and zero-shot from Q2.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e51881d6-b641-4c52-913a-0d14182b40db",
   "metadata": {},
   "outputs": [],
   "source": [
    "def classify_few_shot(train_data, test_data):\n",
    "    \"\"\"Classification with 3-shots.\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "488c8412-3382-403a-87fc-0d4e6ef871d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def classify_zero_shot_cot(train_data, test_data):\n",
    "    \"\"\"Classification with zero-shot chain-of-thought.\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fb8c954-def4-4e0c-a8e2-0294fad2ea7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def classify_few_shot_cot(train_data, test_data):\n",
    "    \"\"\"Classification with 3-shot chain-of-thought.\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "771d1c6d-9f4b-4751-aec2-65dd22c63465",
   "metadata": {},
   "outputs": [],
   "source": [
    "def classify_zero_shot_self_consistency(train_data, test_data):\n",
    "    \"\"\"Implement self-consistency for zero-shot prompting.\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd8a6bdf-3896-4f32-a23d-5b68d76578de",
   "metadata": {},
   "outputs": [],
   "source": [
    "for name, fn in [\n",
    "    (\"majority\", classify_majority_label),\n",
    "    (\"zero-shot\", classify_zero_shot),\n",
    "    (\"few-shot\", classify_few_shot),\n",
    "    (\"zero-shot-cot\", classify_zero_shot_cot),\n",
    "    (\"few-shot-cot\", classify_few_shot_cot),\n",
    "    (\"self-consistency\", classify_zero_shot_self_consistency)\n",
    "]:\n",
    "    score = evaluate(fn)\n",
    "    print(f\"{name}\\t{score}\")"
   ]
  }
 ],
 "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.11.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}