LLM_Study.ipynb

{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"source":"<a href=\"https://www.kaggle.com/code/kuchiriel/llm-study-fast-inference-dynamic-hardware-and-llms?scriptVersionId=159668091\" target=\"_blank\"><img align=\"left\" alt=\"Kaggle\" title=\"Open in Kaggle\" src=\"https://kaggle.com/static/images/open-in-kaggle.svg\"></a>","metadata":{},"cell_type":"markdown"},{"cell_type":"code","source":"! pip install einops\n! pip install torch\n! pip install torchvision\n! pip install google-api-python-client\n#! pip install torch-xla\n! pip install attention-sinks\n! pip install git+https://github.com/huggingface/transformers\n! pip install bitsandbytes\n! pip install accelerate\n! pip install tensorflow\n! pip install numpy\n! pip install SciPy\n! pip install collections nltk re\n# ! pip install ctransformers\n# ! pip install ctransformers[cuda]\n# ! pip install --upgrade scipy\n# ! pip install datasets\n# ! pip install gradio","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#! pip install einops\n#! pip install torch\n#! pip install torchvision\n#! pip install google-api-python-client\n#! pip install attention-sinks\n#! pip install git+https://github.com/huggingface/transformers\n#! pip install bitsandbytes\n#! pip install accelerate\n#! pip install tensorflow\n#! pip install numpy\n#! pip install SciPy\n#! pip install collections nltk re\n##! pip install torch-xla\n## ! pip install ctransformers\n## ! pip install ctransformers[cuda]\n## ! pip install --upgrade scipy\n## ! pip install datasets\n## ! pip install gradio\n\n# Import Garbage Collector module\nimport gc\n\n# Import os module\nimport os\n\nimport re\n\n# Import the AutoTokenizer class from the attention_sinks module\n# TL;DR: attention_sinks adapts pre-trained LLMs to use a modified\n# form of sliding window attention that remains able to produce fluent text indefinitely.\n# https://github.com/tomaarsen/attention_sinks\nfrom attention_sinks import AutoModelForCausalLM, AutoTokenizer\n\n# Import the GenerationConfig and TextStreamer class from the transformers module\nfrom transformers import GenerationConfig, TextStreamer\n\n# Import the Accelerator class from the accelerate module\nfrom accelerate import Accelerator\n\n# Import the torch library\nimport torch\n\n# Import the nn module from the torch library\n#from torch import nn\n\n# Import tensorflow module\nimport tensorflow as tf\n\nimport nltk\n\nfrom nltk.corpus import stopwords\n\n# import torch_xla.core.xla_model as xm","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"nltk.download(\"stopwords\")\n\nos.environ[\"PYTORCH_CUDA_ALLOC_CONF\"] = \"max_split_size_mb:50\"\n\naccelerator = Accelerator()\nCONTENT = \"You are an assistant that follow user instructions precisely\"\nTOKENIZER_MODEL_MAX_LENGHT = 256\nSAVE_DIR = \"~/.cache/huggingface/transformers\"\n\n\ndef test_is_tpu_available():\n    \"\"\"\n    Checks if a TPU device is available.\n\n    Returns:\n        bool: True if a TPU device is available, False otherwise.\n    \"\"\"\n    devices = tf.config.list_logical_devices()\n    for device in devices:\n        if device.device_type == \"TPU\":\n            return True\n    return False  # Add a return statement here to handle the case when no TPU device is found\n\n\nif torch.cuda.is_available():\n    DEVICE = \"cuda\"\n    print(\"GPU Available:\", torch.cuda.get_device_name(torch.cuda.current_device()))\nelif test_is_tpu_available():\n    try:\n        # Attempt to initialize the TPU\n        tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\n        tf.config.experimental_connect_to_cluster(tpu)\n        tf.tpu.experimental.initialize_tpu_system(tpu)\n        print(f\"Running on TPU: {tpu.cluster_spec().as_dict()['worker']}\")\n        strategy = tf.distribute.experimental.TPUStrategy(tpu)\n        DEVICE = \"tpu\"\n    except tf.errors.AlreadyExistsError:\n        print(\"TPU already initialized\")\n    except tf.errors.FailedPreconditionError as e:\n        print(f\"Failed to initialize TPU: {e}\")\n        DEVICE = \"cpu\"\nelse:\n    DEVICE = \"cpu\"\nprint(\"Device:\", DEVICE)\n\nif DEVICE in [\"cuda\", \"tpu\"]:\n    MODEL = \"HuggingFaceH4/zephyr-7b-beta\"\nelse:\n    MODEL = \"microsoft/phi-2\"\nif DEVICE == \"cuda\" and torch.cuda.is_bf16_supported():\n    DTYPE = torch.bfloat16\nelif DEVICE == \"cuda\" and not torch.cuda.is_bf16_supported():\n    DTYPE = torch.float16\nelse:\n    DTYPE = torch.float32\nprint(DTYPE)\n\n\ndef clear_all():\n    \"\"\"\n    Clears all the global variables and optionally clears the local variables as well.\n\n    Parameters:\n        None\n\n    Returns:\n        None\n    \"\"\"\n    flush()\n    # Optionally, clear the locals too\n    locals_to_remove = [var for var in locals() if var[0] != \"_\"]\n    for var in locals_to_remove:\n        del locals()[var]\n    gc.collect()\n\n    # Get a list of all global variables\n    globals_to_remove = [var for var in globals() if var[0] != \"_\"]\n    for var in globals_to_remove:\n        del globals()[var]\n\n\ndef flush():\n    \"\"\"\n    Flushes the memory by freeing up the allocated memory. If the device is set to \"cuda\" or \"tpu\",\n    it also clears the GPU memory and resets the peak memory stats.\n\n    Parameters:\n        None\n\n    Returns:\n        None\n    \"\"\"\n    accelerator.free_memory()\n    if DEVICE in \"cuda\" or DEVICE in \"tpu\":\n        torch.cuda.empty_cache()\n        torch.cuda.reset_peak_memory_stats()\n\n\ndef print_memory_usage():\n    \"\"\"\n    Print the memory usage of the current CUDA device.\n\n    Parameters:\n        None\n\n    Returns:\n        None\n    \"\"\"\n    if DEVICE == \"cuda\":\n        allocated = torch.cuda.memory_allocated() / 1e9\n        max_allocated = torch.cuda.max_memory_allocated() / 1e9\n        print(f\"Memory Allocated: {allocated} GB, Max Allocated: {max_allocated} GB\")\n        flush()\n\n\ndef check_if_quantized(model):\n    \"\"\"\n    Check if any parameter in the given model is quantized.\n\n    Parameters:\n        model (torch.nn.Module): The model to check.\n\n    Returns:\n        bool: True if any parameter in the model is quantized, False otherwise.\n    \"\"\"\n    for name, param in model.named_parameters():\n        if \"quantized\" in str(param.dtype):\n            print(f\"Layer {name} is quantized with data type: {param.dtype}\")\n            return True\n    return False\n\n\ndef setup_model_and_tokenizer():\n    \"\"\"\n\n    Returns:\n        model (AutoModelForCausalLM): The pretrained language model for text generation.\n        tokenizer (AutoTokenizer): The tokenizer for the language model.\n        generation_config (GenerationConfig): The configuration for text generation.\n    \"\"\"\n    torch.set_grad_enabled(False)\n    model = AutoModelForCausalLM.from_pretrained(\n        MODEL,\n        torch_dtype=DTYPE,\n        attention_sink_size=4,\n        attention_sink_window_size=252,\n        trust_remote_code=True if DEVICE in \"cpu\" else False,\n    )\n\n    tokenizer = AutoTokenizer.from_pretrained(MODEL, use_fast=True)\n\n    model.config.update(\n        {\n            \"load_in_4bit\": check_if_quantized(model),\n            \"pad_token_id\": tokenizer.pad_token_id,\n            \"eos_token_id\": tokenizer.eos_token_id,\n            \"bos_token_id\": tokenizer.bos_token_id,\n            \"unk_token_id\": tokenizer.unk_token_id,\n            \"sep_token_id\": tokenizer.sep_token_id,\n            \"cls_token_id\": tokenizer.cls_token_id,\n            \"mask_token_id\": tokenizer.mask_token_id,\n        }\n    )\n\n    # if torch.cuda.device_count() > 1:\n    #     print(\n    #         f\"{torch.cuda.device_count()} GPUs detected, initializing Data Parallel...\"\n    #     )\n    #     model = nn.DataParallel(model)\n    #     model.to(DTYPE)\n\n    #     underlying_model = (\n    #         model.module if isinstance(model, torch.nn.DataParallel) else model\n    #     )\n\n    #     if underlying_model.config.vocab_size != len(tokenizer):\n    #         underlying_model.resize_token_embeddings(len(tokenizer))\n\n    #     model = underlying_model\n\n    # elif torch.cuda.device_count() < 1:\n    #     model = torch.quantization.quantize_dynamic(\n    #         model, {torch.nn.Linear}, dtype=DTYPE\n    #     )\n\n    generation_config = GenerationConfig(\n        min_length=25,\n        max_length=250,\n        bnb_4bit_compute_dtype=DTYPE,\n        penalty_alpha=0.6,\n        repetition_penalty=1.1,\n        top_k=20,\n        pad_token_id=tokenizer.pad_token_id,\n        eos_token_id=tokenizer.eos_token_id,\n        bos_token_id=tokenizer.bos_token_id,\n        unk_token_id=tokenizer.unk_token_id,\n        sep_token_id=tokenizer.sep_token_id,\n        cls_token_id=tokenizer.cls_token_id,\n    )\n\n    if DEVICE == \"cuda\" and not check_if_quantized(model):\n        model = model.half()\n\n    model = accelerator.prepare(model)\n\n    model.eval()\n\n    print_memory_usage()\n\n    return model, tokenizer, generation_config\n\n\nmodel, tokenizer, generation_config = setup_model_and_tokenizer()\n\n\ndef determine_max_tokens(prompt, base_multiplier=5):\n    \"\"\"\n    Dynamically determines the max_new_tokens for text generation based on the prompt.\n\n    :param prompt: The user's input prompt.\n    :param base_multiplier: Base multiplier to adjust the response length.\n    :return: An integer representing the max_new_tokens value.\n    \"\"\"\n\n    prompt_length = len(tokenizer.encode(prompt))\n\n    # Adjust min and max tokens based on prompt length\n    if prompt_length < 50:  # Short prompt\n        min_tokens, max_tokens = 50, 500\n    elif prompt_length < 100:  # Medium prompt\n        min_tokens, max_tokens = 100, 750\n    else:  # Long prompt\n        min_tokens, max_tokens = 150, 1000\n\n    # Further adjustments based on prompt type or complexity\n    if is_complex(prompt):\n        max_tokens += 200  # Allow longer responses for complex prompts\n\n    if prompt.strip().endswith(\"?\"):\n        multiplier = base_multiplier * 1.5\n    else:\n        multiplier = base_multiplier\n\n    # Calculate the maximum number of new tokens to generate\n    calculated_max_tokens = int(prompt_length * multiplier)\n\n    # Ensure the calculated max tokens is within the defined range\n    max_tokens = max(min_tokens, min(calculated_max_tokens, max_tokens))\n\n    return max_tokens, min_tokens\n\n\ndef is_complex(prompt):\n    \"\"\"\n    Determine if a prompt is complex based on various linguistic features.\n\n    :param prompt: A string representing the prompt to be evaluated.\n    :return: A boolean value indicating whether the prompt is complex.\n    \"\"\"\n    # Define thresholds for complexity\n    word_count_threshold = 12\n    unique_word_threshold = 10\n    long_word_threshold = 7\n    complex_sentence_threshold = 2\n\n    # Tokenize the prompt into words\n    words = prompt.split()\n\n    # Check for word count\n    if len(words) > word_count_threshold:\n        return True\n\n    # Check for unique word count\n    unique_words = set(words)\n    if len(unique_words) > unique_word_threshold:\n        return True\n\n    # Check for long words (indicative of advanced vocabulary)\n    long_words = [word for word in words if len(word) >= long_word_threshold]\n    if len(long_words) > long_word_threshold:\n        return True\n\n    # Check for complex sentence structure\n    sentences = re.split(r\"[.!?]+\", prompt)\n    complex_sentences = [\n        sentence\n        for sentence in sentences\n        if len(sentence.split()) > word_count_threshold\n    ]\n    if len(complex_sentences) > complex_sentence_threshold:\n        return True\n\n    # Check for low-frequency or specialized words (excluding common stopwords)\n    stop_words = set(stopwords.words(\"english\"))\n    non_stop_words = [word for word in unique_words if word.lower() not in stop_words]\n    if len(non_stop_words) > unique_word_threshold:\n        return True\n\n    return False\n\n\ndef stream_text(user_prompt, model, tokenizer, generation_config):\n    # If the model is wrapped by DataParallel, access the original model\n    if isinstance(model, torch.nn.DataParallel):\n        model = model.module\n    \"\"\"\n    Generates a stream of text based on the user prompt.\n\n    Parameters:\n        user_prompt (str): The user prompt to generate text from.\n        model: The language model used for text generation.\n        tokenizer: The tokenizer used to encode the user prompt.\n        generation_config: The configuration for text generation.\n\n    Returns:\n        str: The generated text based on the user prompt.\n    \"\"\"\n    messages = [\n        {\n            \"role\": \"system\",\n            \"content\": CONTENT,\n        },\n        {\"role\": \"user\", \"content\": user_prompt},\n    ]\n    user_prompt = tokenizer.apply_chat_template(\n        messages, tokenize=False, return_tensors=\"pt\"\n    )\n    input_ids = tokenizer.encode(user_prompt, return_tensors=\"pt\").to(model.device)\n    max_tokens, min_tokens = determine_max_tokens(user_prompt)\n\n    streamer = TextStreamer(tokenizer)\n    generated_tokens = model.generate(\n        input_ids,\n        min_length=min_tokens,\n        max_length=max_tokens,\n        generation_config=generation_config,\n        streamer=streamer,\n    )\n\n    output_text = tokenizer.decode(generated_tokens[0], skip_special_tokens=True)\n    return output_text\n\n\ndef call_llm():\n    \"\"\"\n    Executes the call to the Language Learning Model (LLM).\n\n    This function prompts the user to enter a prompt, which will be used as input\n    to the LLM for text generation. The user can exit the function by typing\n    'exit'.\n\n    Parameters:\n        None\n\n    Returns:\n        str: The user prompt entered by the user.\n\n    Raises:\n        Exception: If an error occurs during text generation.\n\n    Examples:\n        >>> call_llm()\n        Enter your prompt (or type 'exit' to quit): What is the meaning of life?\n        'What is the meaning of life?'\n    \"\"\"\n    try:\n        user_prompt = input(\"Enter your prompt (or type 'exit' to quit): \")\n        # user_prompt = \"Question: \" + user_prompt + \"\\n\\nAnswer:\"\n        # We use the tokenizer's chat template to format each message -\n        # see https://huggingface.co/docs/transformers/main/en/chat_templatin\n        if user_prompt.lower() == \"exit\":\n            return \"exit\"\n        if DEVICE == \"cuda\":\n            with torch.inference_mode():\n                stream_text(user_prompt, model, tokenizer, generation_config)\n        else:\n            stream_text(user_prompt, model, tokenizer, generation_config)\n    except Exception as ex:\n        print(f\"An error occurred during text generation: {ex}\")\n    finally:\n        if DEVICE == \"cuda\":\n            print_memory_usage()\n        else:\n            accelerator.free_memory()\n    return user_prompt\n\n\ndef loop_llm():\n    \"\"\"\n    A function that continuously loops until a specific condition is met or the program is interrupted.\n    This function calls the `call_llm()` function in a loop and checks the return value. If the return\n    value is \"exit\", it performs some cleanup tasks including printing memory usage, deleting the `model`\n    and `tokenizer` objects, and running garbage collection. Finally, it breaks out of the loop and\n    terminates the function. If the program is interrupted by a `KeyboardInterrupt`, it prints a message\n    and exits gracefully.\n\n    Parameters:\n        None\n\n    Return:\n        None\n    \"\"\"\n    try:\n        while True:\n            if call_llm() == \"exit\":\n                break\n    except KeyboardInterrupt:\n        print(\"\\nExiting the program.\")\n        print_memory_usage()\n        clear_all()\n\n\nloop_llm()\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}