text_completion.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama import Llama\n",
    "from typing import List\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ckpt_dir = \"Meta-Llama-3-8B\"\n",
    "tokenizer_path = \"Meta-Llama-3-8B/tokenizer.model\"\n",
    "temperature: float = 0.6\n",
    "top_p: float = 0.9\n",
    "max_seq_len: int = 128\n",
    "max_gen_len: int = 64\n",
    "max_batch_size: int = 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import torch.distributed as dist\n",
    "os.environ['MASTER_ADDR'] = 'localhost'\n",
    "os.environ['MASTER_PORT'] = '12355'\n",
    "os.environ['RANK'] = '0'\n",
    "os.environ['WORLD_SIZE'] = '1'\n",
    "dist.init_process_group(backend='gloo', init_method='env://')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "generator = Llama.build(\n",
    "        ckpt_dir=ckpt_dir,\n",
    "        tokenizer_path=tokenizer_path,\n",
    "        max_seq_len=max_seq_len,\n",
    "        max_batch_size=max_batch_size,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompts: List[str] = [\n",
    "        # For these prompts, the expected answer is the natural continuation of the prompt\n",
    "        \"I believe the meaning of life is\",\n",
    "        \"Simply put, the theory of relativity states that \",\n",
    "        \"\"\"A brief message congratulating the team on the launch:\n",
    "\n",
    "        Hi everyone,\n",
    "\n",
    "        I just \"\"\",\n",
    "        # Few shot prompt (providing a few examples before asking model to complete more);\n",
    "        \"\"\"Translate English to French:\n",
    "\n",
    "        sea otter => loutre de mer\n",
    "        peppermint => menthe poivrée\n",
    "        plush girafe => girafe peluche\n",
    "        cheese =>\"\"\",\n",
    "    ]\n",
    "results = generator.text_completion(\n",
    "        prompts,\n",
    "        max_gen_len=max_gen_len,\n",
    "        temperature=temperature,\n",
    "        top_p=top_p,\n",
    "    )\n",
    "\n",
    "for prompt, result in zip(prompts, results):\n",
    "    print(prompt)\n",
    "    print(f\"> {result['generation']}\")\n",
    "    print(\"\\n==================================\\n\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "fastai",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}