{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8ada5e87-9fec-4872-8f32-986e9633f574",
   "metadata": {},
   "outputs": [],
   "source": [
    "# IMPORTS \n",
    "import torch\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "02787ef8-d1bf-45a4-b514-441ffa60e742",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"infly/OpenCoder-1.5B-Instruct\"\n",
    "model = AutoModelForCausalLM.from_pretrained(model_name,\n",
    "                                                        torch_dtype = torch.bfloat16,\n",
    "                                                         device_map=\"auto\",\n",
    "                                                         trust_remote_code=True)\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "01e33504-9263-4b0d-a122-3d6cd0ec6de5",
   "metadata": {},
   "outputs": [],
   "source": [
    "messages=[\n",
    "   { 'role': 'user', 'content': 'write a quick sort algorithm in python'}  \n",
    "]\n",
    "\n",
    "inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors=\"pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "49439016-e992-40d6-b3ed-58558d71abe9",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
      "Setting `pad_token_id` to `eos_token_id`:None for open-end generation.\n",
      "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": [
      "Here is a simple implementation of the quick sort algorithm in Python:\n",
      "\n",
      "```python\n",
      "def quick_sort(arr):\n",
      "    if len(arr) <= 1:\n",
      "        return arr\n",
      "    else:\n",
      "        pivot = arr[len(arr) // 2]\n",
      "        left = [x for x in arr if x < pivot]\n",
      "        middle = [x for x in arr if x == pivot]\n",
      "        right = [x for x in arr if x > pivot]\n",
      "        return quick_sort(left) + middle + quick_sort(right)\n",
      "\n",
      "print(quick_sort([3,6,8,10,1,2,1]))\n",
      "```\n",
      "\n",
      "This quick sort algorithm works by selecting a 'pivot' element from the array and partitioning the other elements into two sub-arrays, according to whether they are less than or greater than the pivot. The sub-arrays are then recursively sorted.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "outputs = model.generate(inputs, max_new_tokens=512, do_sample=False)\n",
    "\n",
    "result = tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True)\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5b4a28b-664d-4329-832e-4f1a684b4bd7",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}