LLM-Llama.ipynb

llm-llama.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "machine_shape": "hm",
      "gpuType": "A100",
      "mount_file_id": "1G9kvKC4r7rKhosS8jIIiNUB-0D6EwUp-",
      "authorship_tag": "ABX9TyNdWoki18s3rfy9arOHn7AI",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "56cb6bc98c114fc48f1f47a454590317": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_425b9c3812914ec1a4a9ef520cdd0fbf",
              "IPY_MODEL_c7d685707c6f4da6bfc89551aa35455f",
              "IPY_MODEL_635f5ee14eca489186363f696c05a178"
            ],
            "layout": "IPY_MODEL_7fd4134299a9428ab488dac19d58d6aa"
          }
        },
        "425b9c3812914ec1a4a9ef520cdd0fbf": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_14ff18f6079d4333b515ea37f3bf3aa0",
            "placeholder": "​",
            "style": "IPY_MODEL_37151c6b9d12447abef15cdd3bfcbd85",
            "value": "Loading checkpoint shards: 100%"
          }
        },
        "c7d685707c6f4da6bfc89551aa35455f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_2ffbccb8959542a9a2f7ca34e19d33e0",
            "max": 2,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_f76f7d17abf84002ac64227335a3805c",
            "value": 2
          }
        },
        "635f5ee14eca489186363f696c05a178": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_8b8d7872f9c74c5eb5e01ee608f6302b",
            "placeholder": "​",
            "style": "IPY_MODEL_57e4c69bf785475494a40c29d3f22ab9",
            "value": " 2/2 [01:00<00:00, 27.68s/it]"
          }
        },
        "7fd4134299a9428ab488dac19d58d6aa": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "14ff18f6079d4333b515ea37f3bf3aa0": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "37151c6b9d12447abef15cdd3bfcbd85": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "2ffbccb8959542a9a2f7ca34e19d33e0": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "f76f7d17abf84002ac64227335a3805c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "8b8d7872f9c74c5eb5e01ee608f6302b": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "57e4c69bf785475494a40c29d3f22ab9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/gist/ZhiyaoShu/693e63ad6967f8e67bce375a908233fa/llm-llama.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install transformers einops accelerate langchain bitsandbytes"
      ],
      "metadata": {
        "id": "sS5OHvzuwqlV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install tiktoken"
      ],
      "metadata": {
        "id": "CeTpiVD2x3s2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install accelerate"
      ],
      "metadata": {
        "id": "nOHEgvJs2byu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tVIvThh2v5C-"
      },
      "outputs": [],
      "source": [
        "import transformers\n",
        "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
        "import json\n",
        "from langchain import HuggingFacePipeline\n",
        "from transformers import pipeline\n",
        "from transformers import AutoTokenizer, LlamaTokenizer, LlamaForCausalLM\n",
        "from huggingface_hub import notebook_login\n",
        "import pandas as pd\n",
        "import os\n",
        "import torch\n",
        "from accelerate import Accelerator\n",
        "import tiktoken\n",
        "import logging\n",
        "from tqdm import tqdm\n",
        "import numpy as np\n",
        "import torch\n",
        "import pandas as pd\n",
        "import gc\n",
        "huggingface_token = os.environ.get(\"YOUR_HUGGING_FACE_KEY\")\n",
        "\n",
        "notebook_login(huggingface_token)\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Clean cache to save spaces\n",
        "gc.collect()\n",
        "torch.cuda.empty_cache()"
      ],
      "metadata": {
        "id": "PTuwVjWD-QZ2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Initialize model through hugging-face token\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"meta-llama/Llama-2-7b-chat-hf\")\n",
        "tokenizer.pad_token = tokenizer.eos_token\n",
        "tokenizer.padding_side = \"right\"\n",
        "model_llama = LlamaForCausalLM.from_pretrained(\n",
        "   \"meta-llama/Llama-2-7b-chat-hf\",\n",
        "    torch_dtype=torch.float16,\n",
        "    eos_token_id=tokenizer.eos_token_id,\n",
        ")\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 49,
          "referenced_widgets": [
            "56cb6bc98c114fc48f1f47a454590317",
            "425b9c3812914ec1a4a9ef520cdd0fbf",
            "c7d685707c6f4da6bfc89551aa35455f",
            "635f5ee14eca489186363f696c05a178",
            "7fd4134299a9428ab488dac19d58d6aa",
            "14ff18f6079d4333b515ea37f3bf3aa0",
            "37151c6b9d12447abef15cdd3bfcbd85",
            "2ffbccb8959542a9a2f7ca34e19d33e0",
            "f76f7d17abf84002ac64227335a3805c",
            "8b8d7872f9c74c5eb5e01ee608f6302b",
            "57e4c69bf785475494a40c29d3f22ab9"
          ]
        },
        "id": "fVGeO1HhxmPt",
        "outputId": "d82b14c0-8298-499b-9162-71b0a95bc212"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "56cb6bc98c114fc48f1f47a454590317"
            }
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Setup logging\n",
        "logging.basicConfig(level=logging.INFO)\n",
        "logger = logging.getLogger(__name__)\n",
        "\n",
        "# Example function to query a model with error handling\n",
        "def query_llama(prompt, model_name=\"meta-llama/Llama-2-7b-chat-hf\"):\n",
        "    try:\n",
        "        # Initialize the pipeline with the specified model and tokenizer\n",
        "        pipe = pipeline(\n",
        "            task=\"text-generation\",\n",
        "            model=\"meta-llama/Llama-2-7b-chat-hf\",\n",
        "            tokenizer=tokenizer,\n",
        "            device=0,\n",
        "            return_full_text=False,\n",
        "        )\n",
        "        # Generate text\n",
        "        result = pipe(prompt, truncation=True)\n",
        "        output_text = result[0][\"generated_text\"]\n",
        "        return output_text\n",
        "\n",
        "    except Exception as e:\n",
        "        logger.error(f\"Failed to query LLAMA API: {str(e)}\")\n",
        "        return None\n"
      ],
      "metadata": {
        "id": "Rs5suKWoxfJF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Example usage to test\n",
        "prompt = \"Write me a poem about Machine Learning.\"\n",
        "response = query_llama(prompt)\n",
        "\n",
        "if response:\n",
        "    print(response)\n",
        "else:\n",
        "    print(\"Failed to generate text.\")"
      ],
      "metadata": {
        "id": "lZbNM-1S0nhQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Utilize customized\n",
        "\n",
        "# INSTRUCTION = \"\"\"\n",
        "# Based on the following descriptions, generate one probable MBTI personality type for the corresponding user. IMPORTANT: NO OTHER TEXTS, EXPLANATIONS OR COMMENTARY ARE INCLUDED IN ANY CIRCUMSTANCES! ONLY GENERATE 4 LETTERS OF THE MBTI TYPE, LIKE: \"INFJ\",\"INFP\",\"ENFJ\",\"ENFP\", etc. If you are unsure about a particular user based on the description, or you don't have an accurate prediction, you still have to give one assumption of the MBTI type that you think is the most likely. Your analysis should be based solely on the provided descriptions.\n",
        "# \"\"\"\n",
        "\n",
        "def load_json(file_path):\n",
        "    try:\n",
        "        with open(file_path, \"r\") as f:\n",
        "            data = json.load(f)\n",
        "        return data\n",
        "    except Exception as e:\n",
        "        print(\"Failed to read descriptions\")\n",
        "        raise e\n",
        "\n",
        "def predict_mbti(file_path, model_function):\n",
        "    \"\"\"\n",
        "    Processes the descriptions and predicts MBTI using the specified model.\n",
        "\n",
        "    :param file_path: Path to the JSON file containing user descriptions.\n",
        "    :param model_name: Name of the model to use for prediction.\n",
        "    :param query_function: The function to use for querying the model.\n",
        "    :return: The prediction result or None if an error occurs.\n",
        "    \"\"\"\n",
        "    try:\n",
        "        with open(file_path, 'r', encoding='utf-8') as file:\n",
        "            descriptions = json.load(file)\n",
        "    except Exception as e:\n",
        "        print(f\"Failed to process the descriptions with {model_function.__name__}: {e}\")\n",
        "        return\n",
        "    results = {}\n",
        "\n",
        "    for user_id, description in tqdm(descriptions.items()):\n",
        "      prompt = f\"Based on the provided {description}, generate one most probable MBTI personality type for the user. This response MUST be in the format consisting STRICTLY of one of the 16 MBTI types (e.g., INFP, ESTJ) WITHOUT descriptions, texts, explanations, or commentary.\"\n",
        "\n",
        "      try:\n",
        "          prediction = model_function(prompt)\n",
        "          results[user_id] = prediction\n",
        "\n",
        "      except Exception as e:\n",
        "            print(f\"Failed to process user {user_id} with {model_function.__name__}: {e}\")\n",
        "    return results\n",
        "\n",
        "def generate_and_save_results(file_path):\n",
        "    \"\"\"\n",
        "    Generates MBTI predictions using all three models and saves the results to separate files.\n",
        "\n",
        "    :param file_path: Path to the JSON file containing user descriptions.\n",
        "    \"\"\"\n",
        "    model_functions = {\n",
        "        'llama': query_llama,\n",
        "    }\n",
        "\n",
        "    for model_name in tqdm(model_functions.keys(), desc=\"Processing models\"):\n",
        "        model_function = model_functions[model_name]\n",
        "        output_filename = f\"mbti_predictions_{model_name}.json\"\n",
        "        try:\n",
        "            response = predict_mbti(file_path, model_function)\n",
        "            if response:\n",
        "                save_to_json(response, output_filename)\n",
        "                print(f\"Successfully saved predictions from {model_name} to {output_filename}\")\n",
        "            else:\n",
        "                print(f\"Failed to generate MBTI predictions using {model_name}.\")\n",
        "        except Exception as e:\n",
        "            print(f\"Exception while processing {model_name}: {e}\")\n",
        "\n",
        "def save_to_json(data, filename):\n",
        "    try:\n",
        "        with open(\"/content/drive/MyDrive/000/prediction_llm.json\", \"w\") as f:\n",
        "            json.dump(data, f, ensure_ascii=False, indent=4)\n",
        "        print(f\"Data successfully saved to {filename}\")\n",
        "    except Exception as e:\n",
        "        print(f\"Failed to save data to {filename}: {e}\")\n",
        "\n",
        "\n",
        "file_path = \"/content/drive/MyDrive/000/gpt_description.json\"\n",
        "with torch.no_grad():\n",
        "    generate_and_save_results(file_path)\n",
        "    pass"
      ],
      "metadata": {
        "id": "xSLMxLUZ6jm4"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}