LLMの脳波をとってみる

analyse_llm.py

# Dependencies: transformers, colorama, torch, matplotlib, numpy
# pip install transformers colorama torch matplotlib numpy
# Models are stored at C:\Users\%username%\.cache\huggingface\hub
import time
import random
import torch
import array
import numpy as np
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, Conversation, set_seed
from colorama import Fore, Back, Style, init
import matplotlib.pyplot as plt
init(autoreset=True)

max_length=256

#model_name = "cerebras/Cerebras-GPT-111M"
model_name = "cerebras/Cerebras-GPT-256M"
range_max = 125
#model_name = "cerebras/Cerebras-GPT-590M"
#model_name = "cerebras/Cerebras-GPT-1.3B"
#model_name = "cerebras/Cerebras-GPT-2.7B"
#model_name = "cerebras/Cerebras-GPT-6.7B"
#model_name = "cerebras/Cerebras-GPT-13B"
#model_name = "rinna/japanese-gpt2-medium"
#range_max = 40

def banner(model_name, max_length):

    print("Model : " + Fore.RED + f"{model_name}")
    print("Params: " + Fore.YELLOW + f"max_new_tokens:{max_length}")


def analyse_gpt(model_name, max_length):

    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name)
    model.config.pad_token_id = model.config.eos_token_id
    print("module size: ", len(model.transformer.h))

    cmap = plt.cm.viridis

    summary = []
    for idx, elm in enumerate(model.transformer.h):
        elm.ln_1.index = idx * 2
        elm.ln_2.index = idx * 2 + 1
        elm.ln_1.old_forward = elm.ln_1.forward
        elm.ln_2.old_forward = elm.ln_2.forward
        def new_forward(self, x):
            result = self.old_forward(x)
            ar = result.detach().numpy()
            summary[self.index] += ar[0][0] ** 2
            return result
        elm.ln_1.forward = new_forward.__get__(elm.ln_1)
        elm.ln_2.forward = new_forward.__get__(elm.ln_2)

        tensor1 = elm.ln_1.weight.data
        tensor2 = elm.ln_2.weight.data
        summary.append(np.zeros(len(tensor1), dtype=float))
        summary.append(np.zeros(len(tensor2), dtype=float))

    opts = {}
    opts["max_new_tokens"] = max_length

    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)#, device_map='auto')

    while True:
        summary = []
        for elm in model.transformer.h:
            tensor1 = elm.ln_1.weight.data
            tensor2 = elm.ln_2.weight.data
            summary.append(np.zeros(len(tensor1), dtype=float))
            summary.append(np.zeros(len(tensor2), dtype=float))

        prompt = input("prompt: ")
        if prompt == "exit":
            break
        generated_text = pipe(prompt, do_sample=True, use_cache=True, **opts)[0]   
        print(Fore.YELLOW + prompt + Fore.WHITE + generated_text['generated_text'][len(prompt):])     
        plt.imshow(np.clip(np.sqrt(summary), 0, range_max), cmap=cmap, aspect='auto')
        plt.colorbar()
        plt.show()

if __name__ == '__main__':

    banner(model_name, max_length)
    start_time=time.time()
    analyse_gpt(model_name, max_length)
    end_time=time.time()

    print(Style.BRIGHT + Fore.RED + f"Execution time: {end_time - start_time:.2f} sec")