{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "CIUy2FYzfFEJ"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "!pip install transformers\n",
        "!pip install einops\n",
        "!pip install accelerate\n",
        "\n",
        "import torch\n",
        "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
        "\n",
        "model = AutoModelForCausalLM.from_pretrained(\"agonh/phi-base_model\", trust_remote_code=True, device_map=\"cuda:0\")\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"agonh/phi-base_model\", trust_remote_code=True, device_map=\"cuda:0\")"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "prompt = \"tell me about the moon ?\"\n",
        "inputs = tokenizer(prompt, return_tensors=\"pt\", return_attention_mask=False)\n",
        "\n",
        "inputs.to(\"cuda:0\")\n",
        "outputs = model.generate(**inputs, max_length=100)\n",
        "text = tokenizer.batch_decode(outputs)[0]\n",
        "print(text)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "87vmiui3jGK7",
        "outputId": "46fb0454-8b29-4a76-cb1c-1d1c7c7c28c2"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tell me about the moon?\n",
            "\n",
            "Teacher: The moon is a natural satellite of the Earth. It is the fifth largest moon in the solar system and is about 238,900 miles away from the Earth.\n",
            "\n",
            "Student: Wow, that's a really big distance!\n",
            "\n",
            "Teacher: Yes, it is. The moon has a lot of interesting features, like craters and mountains. It also affects the tides on Earth.\n",
            "\n",
            "Student: How does the moon affect the tides\n"
          ]
        }
      ]
    }
  ]
}