{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "aebdeaaf529549a1ad39ccf838e39e71",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer_config.json:   0%|          | 0.00/26.0 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "26da6488303344ce90900207ae33b6bd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "vocab.json:   0%|          | 0.00/1.04M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1eea77d125a74e00b1f050960dfa42aa",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "merges.txt:   0%|          | 0.00/456k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b704ae3896dc46698ca50b7a29230bde",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer.json:   0%|          | 0.00/1.36M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "42eaaf035fad40f583ee01dac1a86e85",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "config.json:   0%|          | 0.00/665 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from transformers import GPT2LMHeadModel, GPT2Tokenizer\n",
    "\n",
    "model = GPT2LMHeadModel.from_pretrained(\"models/math-gpt2-v0/checkpoint-64\")\n",
    "tokenizer = GPT2Tokenizer.from_pretrained(\"gpt2\")\n",
    "model = model.to(\"cuda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The attention mask is not set and cannot be inferred from input because pad token is same as eos token. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Your math prompt here:\n",
      "\n",
      "$ x = (m^3 + k) + (m^4 + k) + (m^5 + k)\n",
      "\n",
      "Now, we have:\n",
      "\n",
      "$ x = (m^3 +\n"
     ]
    }
   ],
   "source": [
    "def generate_text(prompt, max_length=50):\n",
    "\n",
    "    inputs = tokenizer(prompt, return_tensors=\"pt\").to(\"cuda\")\n",
    "    \n",
    "    outputs = model.generate(\n",
    "        inputs[\"input_ids\"],\n",
    "        max_length=max_length,\n",
    "        pad_token_id=tokenizer.eos_token_id,\n",
    "        num_return_sequences=1,\n",
    "        do_sample=True,\n",
    "        temperature=0.7\n",
    "    )\n",
    "    \n",
    "    return tokenizer.decode(outputs[0], skip_special_tokens=True)\n",
    "\n",
    "prompt = \"Your math prompt here\"\n",
    "generated_text = generate_text(prompt)\n",
    "print(generated_text)"
   ]
  }
 ],
 "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}