ibm-granite-test pipeline.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "HUGGINGFACE_API_KEY = os.environ.get('API_KEY')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "model_id = 'ibm-granite/granite-8b-code-base-128k'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir=\"./models/\", use_fast= True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3db67efaae384cd7b3963a92269f83c9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "LlamaForCausalLM(\n",
       "  (model): LlamaModel(\n",
       "    (embed_tokens): Embedding(49152, 4096, padding_idx=0)\n",
       "    (layers): ModuleList(\n",
       "      (0-35): 36 x LlamaDecoderLayer(\n",
       "        (self_attn): LlamaSdpaAttention(\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=True)\n",
       "          (k_proj): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          (v_proj): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          (o_proj): Linear(in_features=4096, out_features=4096, bias=True)\n",
       "          (rotary_emb): LlamaRotaryEmbedding()\n",
       "        )\n",
       "        (mlp): LlamaMLP(\n",
       "          (gate_proj): Linear(in_features=4096, out_features=14336, bias=True)\n",
       "          (up_proj): Linear(in_features=4096, out_features=14336, bias=True)\n",
       "          (down_proj): Linear(in_features=14336, out_features=4096, bias=True)\n",
       "          (act_fn): SiLU()\n",
       "        )\n",
       "        (input_layernorm): LlamaRMSNorm((4096,), eps=1e-05)\n",
       "        (post_attention_layernorm): LlamaRMSNorm((4096,), eps=1e-05)\n",
       "      )\n",
       "    )\n",
       "    (norm): LlamaRMSNorm((4096,), eps=1e-05)\n",
       "    (rotary_emb): LlamaRotaryEmbedding()\n",
       "  )\n",
       "  (lm_head): Linear(in_features=4096, out_features=49152, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = AutoModelForCausalLM.from_pretrained(model_id, device_map=device, cache_dir=\"./models/\", low_cpu_mem_usage=True)\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "my_pipeline = pipeline(\n",
    "    task='text-generation',\n",
    "    model=model,\n",
    "    tokenizer=tokenizer,\n",
    "    max_length=50,  # Limit the sequence length\n",
    "    num_beams=1,  # Use greedy decoding\n",
    "    do_sample=False,  # Disable sampling\n",
    "    return_full_text=False  # Only return generated text\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_text = [\"def fibonacci():\", \"def greatest_common_factor():\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Truncation was not explicitly activated but `max_length` is provided a specific value, please use `truncation=True` to explicitly truncate examples to max length. Defaulting to 'longest_first' truncation strategy. If you encode pairs of sequences (GLUE-style) with the tokenizer you can select this strategy more precisely by providing a specific strategy to `truncation`.\n",
      "Starting from v4.46, the `logits` model output will have the same type as the model (except at train time, where it will always be FP32)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[[{'generated_text': '\\n    a, b = 0, 1\\n    while True:\\n        yield a\\n        a, b = b, a + b\\n\\ndef main():\\n    for i in range(10):\\n        print(next('}],\n",
       " [{'generated_text': '\\n    \"\"\"\\n    Returns the greatest common factor of two numbers.\\n    \"\"\"\\n    a = int(input(\"Enter the first number: \"))\\n    b = int(input(\"Enter the second number: \"))\\n    if a'}]]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "my_pipeline(input_text)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tesis",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}