spider plot code

Author: divnori

xtra_labs/llm_finetune/draft.py

@@ -7,32 +7,31 @@
 import tensorflow as tf
 import transformers
 
-from utils import run_benchmark
+from utils import run_benchmark, make_spider_plot
 
 # Part 1
 
 # TEXT: overview of LLM lab
 # Load pretrained LLM (medium size model)
 
-model_name = "facebook/opt-1.3b" 
-model = transformers.TFAutoModelForCausalLM.from_pretrained(model_name)
+model_name = "facebook/opt-1.3b"
+# had to load non TF version to run benchmarking code
+model = transformers.AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
 tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
 
 # TEXT: explain tokenizer
 # Include cell for tokenizer inspection
 
 # TEXT: explain how LLMs are trained for next token prediction 
 # Write a function to predict next token
-
-def predict_next_token(probs):
+def predict_next_token(probs, tokenizer):
     new_token = np.random.choice(len(probs), p=probs.numpy())
     print(tokenizer.decode(new_token), end='', flush=True)
     return new_token
 
 # TEXT: explain that next token prediction must be called multiple times for inference
 # Call in loop for autoregressive inference
-
-def generate(start_text, num_steps=20, temp=1.): 
+def generate(start_text, model, tokenizer, num_steps=20, temp=1.): 
     print(start_text, end="")
     x = tokenizer.encode(start_text)
     num_start = len(x)
@@ -42,46 +41,43 @@ def generate(start_text, num_steps=20, temp=1.):
         logits = model(input_tensor).logits
         probs = tf.nn.softmax(logits/temp)[0, -1, :]
 
-        new_token = predict_next_token(probs)
+        new_token = predict_next_token(probs, tokenizer)
         x.append(new_token)
 
     output = tokenizer.decode(x[num_start:])
     return output
 
 # Test autoregressive generation
-
 # while True: 
 #     print("\n\n\n\n\n")
 #     input_text = input("Prompt: ")
-#     output = generate(input_text)
+#     output = generate(input_text, model, tokenizer)
 
 # TEXT: some background on LLM benchmarking
 # Load benchmark dataset and evaluate model
-
 dataset = pd.read_csv("benchmark.csv")
-category_accs_1300m, avg_acc_1300m = run_benchmark(model, tokenizer)
+category_accs_1300m, avg_acc_1300m = run_benchmark(model, tokenizer, dataset)
 
 # TEXT: ask them to make a prediction on how accuracy will be affected by different model sizes
 
 # Benchmark smaller model
-
 model_name_350m = "facebook/opt-350m" 
-model_350m = transformers.TFAutoModelForCausalLM.from_pretrained(model_name_350m)
+model_350m = transformers.AutoModelForCausalLM.from_pretrained(model_name_350m, device_map="auto")
 tokenizer_350m = transformers.AutoTokenizer.from_pretrained(model_350m)
 
-category_accs_350m, avg_acc_350m = run_benchmark(model_350m, tokenizer_350m)
+category_accs_350m, avg_acc_350m = run_benchmark(model_350m, tokenizer_350m, dataset)
 
 # Benchmark larger model
-
 model_name_2700m = "facebook/opt-2.7b" 
-model_2700m = transformers.TFAutoModelForCausalLM.from_pretrained(model_name_2700m)
+model_2700m = transformers.AutoModelForCausalLM.from_pretrained(model_name_2700m, device_map="auto")
 tokenizer_2700m = transformers.AutoTokenizer.from_pretrained(model_2700m)
 
-category_accs_2700m, avg_acc_2700m = run_benchmark(model_2700m, tokenizer_2700m)
+category_accs_2700m, avg_acc_2700m = run_benchmark(model_2700m, tokenizer_2700m, dataset)
 
 # Spider plot
 
-print(category_accs_1300m)
+benchmark_data = {"350M-Model": category_accs_350m, "1300M-Model": category_accs_1300m, "2700M-Model": category_accs_2700m}
+make_spider_plot(benchmark_data)
 
 # Part 2
 

xtra_labs/llm_finetune/spider.png

@@ -1,18 +1,16 @@
 """
 Contains functions that the students will not interface with
 """
-
+import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
 import tensorflow as tf
 import torch
 import torch.nn.functional as F
 from tqdm import tqdm
 
-dataset = pd.read_csv("benchmark.csv")
-
-def run_benchmark(model, tokenizer, few_shot=7, num_steps=500, verbose=False):
-    # device = model.device
+def run_benchmark(model, tokenizer, dataset, few_shot=7, num_steps=500, verbose=False):
+    device = model.device
     dataset["Correct"] = 0.0
 
     # Loop through every question in the benchmark
@@ -32,9 +30,8 @@ def run_benchmark(model, tokenizer, few_shot=7, num_steps=500, verbose=False):
 
             # Run the model 
             with torch.no_grad():
-                # converting to tensorflow tensor for model input
-                x = tokenizer.encode(text, return_tensors="pt")
-                logits = model(x).logits.numpy()
+                x = tokenizer.encode(text, return_tensors="pt").to(device)
+                logits = model(x).logits
                 probs = F.softmax(logits, dim=-1)[0, :-1, :]  # shape: [seq_len-1, vocab_size]
                 y = x[0, 1:]  # shape: [seq_len-1]
 
@@ -60,4 +57,48 @@ def run_benchmark(model, tokenizer, few_shot=7, num_steps=500, verbose=False):
     sorted_accs = accs.sort_values()
     print(sorted_accs)
 
-    return sorted_accs, dataset["Correct"].mean()
+    return accs, dataset["Correct"].mean()
+
+def make_spider_plot(data):
+    """
+    Data is a dictionary where keys are different entities
+    Values are pd Series where series indices are plot labels and series values show performance
+    """
+    colors = ['#1aaf6c', '#429bf4', '#d42cea']
+    i = 0
+    for k,v in data.items():
+        labels = v.index.tolist()
+        values = v.values.tolist()
+        
+        num_vars = len(labels)
+        angles = np.linspace(0, 2 * np.pi, num_vars, endpoint=False).tolist()
+        angles += angles[:1]
+        values += values[:1]
+
+        fig, ax = plt.subplots(figsize=(8,6), subplot_kw=dict(polar=True))
+        
+        ax.plot(angles, values, color=colors[i], linewidth=1, label=k)
+        ax.fill(angles, values, color=colors[i], alpha=0.25)
+
+        i+=1
+
+    ax.set_theta_offset(np.pi / 2)
+    ax.set_theta_direction(-1)
+    ax.set_thetagrids(np.degrees(angles[:-1]), labels)
+    for label, angle in zip(ax.get_xticklabels(), angles):
+        if angle in (0, np.pi):
+            label.set_horizontalalignment('center')
+        elif 0 < angle < np.pi:
+            label.set_horizontalalignment('left')
+        else:
+            label.set_horizontalalignment('right')
+
+    ax.set_ylim(0, 1)
+    ax.set_rlabel_position(180 / num_vars)
+
+    ax.legend(loc='upper right', bbox_to_anchor=(1.3, 1.1))
+
+    plt.savefig("spider.png")
+
+
+