added text models

app.py

@@ -20,7 +20,7 @@ import time
 import threading
 import os
 
-from backbone import extract_features, download_all_models, get_model
+from backbone import extract_features, get_model
 from backbone import MODEL_DICT, LAYER_DICT, RES_DICT
 from ncut_pytorch import NCUT, eigenvector_to_rgb
 
@@ -66,7 +66,7 @@ def compute_ncut(
 ):        
     logging_str = ""
     
-    num_nodes = np.prod(features.shape[:3])
+    num_nodes = np.prod(features.shape[:-1])
     if num_nodes / 2 < num_eig:
         # raise gr.Error("Number of eigenvectors should be less than half the number of nodes.")
         gr.Warning("Number of eigenvectors should be less than half the number of nodes.\n" f"Setting num_eig to {num_nodes // 2 - 1}.")
@@ -100,7 +100,7 @@ def compute_ncut(
     )
     logging_str += f"{embedding_method} time: {time.time() - start:.2f}s\n"
 
-    rgb = rgb.reshape(features.shape[:3] + (3,))
+    rgb = rgb.reshape(features.shape[:-1] + (3,))
     return rgb, logging_str, eigvecs
 
 
@@ -584,7 +584,7 @@ def make_output_images_section():
     return output_gallery
 
 def make_parameters_section():
-    gr.Markdown('### Parameters')
+    gr.Markdown("### Parameters <a style='color: #0044CC;' href='https://ncut-pytorch.readthedocs.io/en/latest/how_to_get_better_segmentation/' target='_blank'>Help</a>")
     from backbone import get_all_model_names
     model_names = get_all_model_names()
     model_dropdown = gr.Dropdown(model_names, label="Backbone", value="DiNO(dino_vitb8)", elem_id="model_name")
@@ -605,6 +605,7 @@ def make_parameters_section():
     model_dropdown.change(fn=change_layer_slider, inputs=model_dropdown, outputs=[layer_slider, node_type_dropdown])
     
     with gr.Accordion("➡️ Click to expand: more parameters", open=False):
+        gr.Markdown("<a href='https://ncut-pytorch.readthedocs.io/en/latest/how_to_get_better_segmentation/' target='_blank'>Docs: How to Get Better Segmentation</a>")
         affinity_focal_gamma_slider = gr.Slider(0.01, 1, step=0.01, label="NCUT: Affinity focal gamma", value=0.5, elem_id="affinity_focal_gamma", info="decrease for shaper segmentation")
         num_sample_ncut_slider = gr.Slider(100, 50000, step=100, label="NCUT: num_sample", value=10000, elem_id="num_sample_ncut", info="Nyström approximation")
         sampling_method_dropdown = gr.Dropdown(["fps", "random"], label="NCUT: Sampling method", value="fps", elem_id="sampling_method", info="Nyström approximation")
@@ -822,11 +823,12 @@ with demo:
         )    
     
     with gr.Tab('Text'): 
-        gr.Markdown('=== under construction ===')
-        gr.Markdown('Please see the [Documentation](https://ncut-pytorch.readthedocs.io/en/latest/gallery_llama3/) for example of NCUT on text input.')
-        gr.Markdown('---')
-        gr.Markdown('![ncut](https://ncut-pytorch.readthedocs.io/en/latest/images/gallery/llama3/llama3_layer_31.jpg)')
-    
+        if USE_HUGGINGFACE_SPACE:
+            from app_text import make_demo
+        else:
+            from draft_gradio_app_text import make_demo
+        make_demo()
+        
     with gr.Tab('Compare Models'): 
         def add_one_model(i_model=1):
             with gr.Column(scale=5, min_width=200) as col:
@@ -897,7 +899,11 @@ with demo:
 
 
 if USE_HUGGINGFACE_SPACE:
+    from backbone import download_all_models
+    threading.Thread(target=download_all_models).start()
+    from backbone_text import download_all_models
     threading.Thread(target=download_all_models).start()
+    
     threading.Thread(target=download_all_datasets).start()
     demo.launch()
 else:

app_text.py

@@ -0,0 +1,304 @@
+# %%
+USE_HUGGINGFACE_SPACE = True
+
+if USE_HUGGINGFACE_SPACE:  # huggingface ZeroGPU, dynamic GPU allocation 
+    try:
+        import spaces
+    except ImportError:
+        USE_HUGGINGFACE_SPACE = False  # run on local machine
+
+import gradio as gr
+
+import torch
+import torch.nn.functional as F
+from PIL import Image
+import numpy as np
+import time
+import threading
+import os
+import matplotlib.pyplot as plt
+import matplotlib.colors as mcolors
+import numpy as np
+
+from ncut_pytorch import NCUT, eigenvector_to_rgb
+
+from backbone_text import MODEL_DICT as TEXT_MODEL_DICT
+from backbone_text import LAYER_DICT as TEXT_LAYER_DICT
+
+def compute_ncut(
+    features,
+    num_eig=100,
+    num_sample_ncut=10000,
+    affinity_focal_gamma=0.3,
+    knn_ncut=10,
+    knn_tsne=10,
+    embedding_method="UMAP",
+    num_sample_tsne=300,
+    perplexity=150,
+    n_neighbors=150,
+    min_dist=0.1,
+    sampling_method="fps",
+    metric="cosine",
+):        
+    logging_str = ""
+    print("running ncut")
+    print(features.shape)
+    num_nodes = np.prod(features.shape[:-1])
+    if num_nodes / 2 < num_eig:
+        # raise gr.Error("Number of eigenvectors should be less than half the number of nodes.")
+        gr.Warning("Number of eigenvectors should be less than half the number of nodes.\n" f"Setting num_eig to {num_nodes // 2 - 1}.")
+        num_eig = num_nodes // 2 - 1
+        logging_str += f"Number of eigenvectors should be less than half the number of nodes.\n" f"Setting num_eig to {num_nodes // 2 - 1}.\n"
+    
+    start = time.time()
+    eigvecs, eigvals = NCUT(
+        num_eig=num_eig,
+        num_sample=num_sample_ncut,
+        device="cuda" if torch.cuda.is_available() else "cpu",
+        affinity_focal_gamma=affinity_focal_gamma,
+        knn=knn_ncut,
+        sample_method=sampling_method,
+        distance=metric,
+        normalize_features=False,
+    ).fit_transform(features.reshape(-1, features.shape[-1]))
+    # print(f"NCUT time: {time.time() - start:.2f}s")
+    logging_str += f"NCUT time: {time.time() - start:.2f}s\n"
+    
+    start = time.time()
+    _, rgb = eigenvector_to_rgb(
+        eigvecs,
+        method=embedding_method,
+        num_sample=num_sample_tsne,
+        perplexity=perplexity,
+        n_neighbors=n_neighbors,
+        min_distance=min_dist,
+        knn=knn_tsne,
+        device="cuda" if torch.cuda.is_available() else "cpu",
+    )
+    logging_str += f"{embedding_method} time: {time.time() - start:.2f}s\n"
+
+    rgb = rgb.reshape(features.shape[:-1] + (3,))
+    return rgb, logging_str, eigvecs
+
+
+def make_plot(token_texts, rgb, num_lines=50, title=""):
+    fig, ax = plt.subplots(figsize=(10, 20))
+    # Define the colors
+    # fill nan with 0
+    rgb = np.nan_to_num(rgb)
+    colors = [mcolors.rgb2hex(rgb[i]) for i in range(len(token_texts))]
+
+    # Split the sentence into words
+    words = token_texts
+
+    y_pos = 0.96
+    x_pos = 0.0
+    max_word_length = max(len(word) for word in words)
+    count = 0
+    for word, color in zip(words, colors):
+        if '\n' in word:
+            word = word.replace('\n', '')
+            y_pos -= 0.025
+            x_pos = 0.0
+            count += 1
+            if count >= num_lines:
+                break
+        
+        text_color = 'black' if sum(mcolors.hex2color(color)) > 1.3 else 'white'  # Choose text color based on background color
+        # text_color = 'black'
+        txt = ax.text(x_pos, y_pos, word, color=text_color, fontsize=12, bbox=dict(facecolor=color, alpha=0.8, edgecolor='none', pad=2))
+        txt_width = txt.get_window_extent().width / (fig.dpi * fig.get_size_inches()[0])  # Calculate the width of the text in inches
+        
+        x_pos += txt_width * 1.2 + 0.01  # Adjust the spacing between words
+        
+        if x_pos > 0.97:
+            y_pos -= 0.025
+            x_pos = 0.0
+            count += 1
+            if count >= num_lines:
+                break
+        # break
+            
+    # Remove the axis ticks and spines
+    ax.set_xticks([])
+    ax.set_yticks([])
+    ax.spines['top'].set_visible(False)
+    ax.spines['right'].set_visible(False)
+    ax.spines['bottom'].set_visible(False)
+    ax.spines['left'].set_visible(False)
+    
+    ax.set_title(title, fontsize=20)
+
+    return fig
+
+
+
+def ncut_run(
+    model,
+    text,
+    model_name,
+    layer=-1,
+    num_eig=100,
+    node_type="block",
+    affinity_focal_gamma=0.3,
+    num_sample_ncut=10000,
+    knn_ncut=10,
+    embedding_method="UMAP",
+    num_sample_tsne=1000,
+    knn_tsne=10,
+    perplexity=500,
+    n_neighbors=500,
+    min_dist=0.1,
+    sampling_method="fps",
+):
+    logging_str = ""
+    if perplexity >= num_sample_tsne or n_neighbors >= num_sample_tsne:
+        # raise gr.Error("Perplexity must be less than the number of samples for t-SNE.")
+        gr.Warning("Perplexity/n_neighbors must be less than the number of samples.\n" f"Setting Perplexity to {num_sample_tsne-1}.")
+        logging_str += f"Perplexity/n_neighbors must be less than the number of samples.\n" f"Setting Perplexity to {num_sample_tsne-1}.\n"
+        perplexity = num_sample_tsne - 1
+        n_neighbors = num_sample_tsne - 1
+
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        
+    node_type = node_type.split(":")[0].strip()
+    
+    model = model.to("cuda" if torch.cuda.is_available() else "cpu")
+    
+    start = time.time()
+    out = model(text)
+    features = out[node_type][layer-1].squeeze(0).detach().float()
+    token_texts = out["token_texts"]
+    
+    if perplexity >= features.shape[0] or n_neighbors >= features.shape[0]:
+        # raise gr.Error("Perplexity must be less than the number of samples.")
+        gr.Warning("Perplexity/n_neighbors must be less than the number of samples.\n" f"Setting Perplexity to {features.shape[0]-1}.")
+        logging_str += f"Perplexity/n_neighbors must be less than the number of samples.\n" f"Setting Perplexity to {features.shape[0]-1}.\n"
+        perplexity = features.shape[0] - 1
+        n_neighbors = features.shape[0] - 1
+    
+    # print(f"Feature extraction time (gpu): {time.time() - start:.2f}s")
+    logging_str += f"Backbone time: {time.time() - start:.2f}s\n"
+    
+    rgb, _logging_str, _ = compute_ncut(
+        features,
+        num_eig=num_eig,
+        num_sample_ncut=num_sample_ncut,
+        affinity_focal_gamma=affinity_focal_gamma,
+        knn_ncut=knn_ncut,
+        knn_tsne=knn_tsne,
+        num_sample_tsne=num_sample_tsne,
+        embedding_method=embedding_method,
+        perplexity=perplexity,
+        n_neighbors=n_neighbors,
+        min_dist=min_dist,
+        sampling_method=sampling_method,
+    )
+    logging_str += _logging_str
+    
+    title = f"{model_name}, Layer {layer}, {node_type}"
+    fig = make_plot(token_texts, rgb, title=title)
+    return fig, logging_str
+
+def _ncut_run(*args, **kwargs):
+    try:
+        ret = ncut_run(*args, **kwargs)
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        return ret
+    except Exception as e:
+        gr.Error(str(e))
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        return None, "Error: " + str(e)
+    
+if USE_HUGGINGFACE_SPACE:
+    @spaces.GPU(duration=30)
+    def __ncut_run(*args, **kwargs):
+        return _ncut_run(*args, **kwargs)
+else:
+    def __ncut_run(*args, **kwargs):
+        return _ncut_run(*args, **kwargs)
+
+def real_run(model_name, text, layer, node_type, num_eig, affinity_focal_gamma, num_sample_ncut, knn_ncut, embedding_method, num_sample_tsne, knn_tsne, perplexity, n_neighbors, min_dist, sampling_method):
+    model = TEXT_MODEL_DICT[model_name]()
+    return __ncut_run(model, text, model_name, layer, num_eig, node_type, 
+                affinity_focal_gamma, num_sample_ncut, knn_ncut, embedding_method, 
+                num_sample_tsne, knn_tsne, perplexity, n_neighbors, min_dist, sampling_method)
+
+lines = \
+"""1. The majestic giraffe, with its towering height and distinctive long neck, roams the savannas of Africa. These gentle giants use their elongated tongues to pluck leaves from the tallest trees, making them well-adapted to their environment. Their unique coat patterns, much like human fingerprints, are unique to each individual.
+2. Penguins, the tuxedoed birds of the Antarctic, are expert swimmers and divers. These flightless seabirds rely on their dense, waterproof feathers and streamlined bodies to propel through icy waters in search of fish, krill, and other marine life. Their huddled colonies and amusing waddles make them a favorite among wildlife enthusiasts.
+3. The mighty African elephant, the largest land mammal, is revered for its intelligence and strong family bonds. These gentle giants use their versatile trunks for various tasks, from drinking and feeding to communicating and greeting one another. Their massive ears and wrinkled skin make them an iconic symbol of the African wilderness.
+4. The colorful and flamboyant peacock, native to Asia, is known for its stunning iridescent plumage. During mating season, the males fan out their magnificent train of feathers, adorned with intricate eye-like patterns, in an elaborate courtship display to attract potential mates, making them a true spectacle of nature.
+5. The sleek and powerful cheetah, the fastest land animal, is built for speed and agility. With its distinctive black tear-like markings and slender body, this feline predator can reach top speeds of up to 70 mph during short bursts, allowing it to chase down its prey with remarkable precision.
+6. The playful and intelligent dolphin, a highly social marine mammal, is known for its friendly demeanor and impressive acrobatic abilities. These aquatic creatures use echolocation to navigate and hunt, and their complex communication systems have long fascinated researchers studying their intricate social structures and cognitive abilities.
+7. The majestic bald eagle, the national emblem of the United States, soars high above with its distinctive white head and tail feathers. These powerful raptors are skilled hunters, swooping down from great heights to catch fish and other prey with their sharp talons, making them an iconic symbol of strength and freedom.
+8. The industrious beaver, nature's skilled engineers, are known for their remarkable ability to construct dams and lodges using their sharp incisors and webbed feet. These semiaquatic rodents play a crucial role in shaping their aquatic ecosystems, creating habitats for numerous other species while demonstrating their ingenuity and perseverance.
+9. The vibrant and enchanting hummingbird, one of the smallest bird species, is a true marvel of nature. With their rapidly flapping wings and ability to hover in mid-air, these tiny feathered creatures are expert pollinators, flitting from flower to flower in search of nectar and playing a vital role in plant reproduction.
+10. The majestic polar bear, the apex predator of the Arctic, is perfectly adapted to its icy environment. With its thick insulating fur and specialized paws for gripping the ice, this powerful carnivore relies on its exceptional hunting skills and keen senses to locate and capture seals, its primary prey, in the harsh Arctic landscape.
+"""
+
+def make_demo():
+    with gr.Row():
+        with gr.Column(scale=5, min_width=200):
+            gr.Markdown("### Input Text")
+            placeholder = lines
+            input_text = gr.Text(value=placeholder, label="Input Text", placeholder="Type here", lines=12)
+            submit_button = gr.Button("🔴 RUN", elem_id="submit_button", variant='primary')
+            clear_button = gr.Button("🗑️Clear", elem_id='clear_button', variant='stop')
+        with gr.Column(scale=5, min_width=200):
+            gr.Markdown("### Parameters <a style='color: #0044CC;' href='https://ncut-pytorch.readthedocs.io/en/latest/how_to_get_better_segmentation/' target='_blank'>Help</a>")
+            model_name = gr.Dropdown(list(TEXT_MODEL_DICT.keys()), label="Model", value="meta-llama/Meta-Llama-3.1-8B")
+            layer = gr.Slider(1, 32, step=1, value=32, label="Layer")
+            node_type = gr.Dropdown(["attn: attention output", "mlp: mlp output", "block: sum of residual"], label="Node Type", value="block: sum of residual")
+            num_eig = gr.Slider(minimum=1, maximum=1000, step=1, value=100, label="Number of Eigenvectors")
+            
+            with gr.Accordion("➡️ Click to expand: more parameters", open=False):
+                gr.Markdown("<a href='https://ncut-pytorch.readthedocs.io/en/latest/how_to_get_better_segmentation/' target='_blank'>Docs: How to Get Better Segmentation</a>")
+                affinity_focal_gamma = gr.Slider(minimum=0.1, maximum=1.0, step=0.1, value=0.3, label="Affinity Focal Gamma")
+                num_sample_ncut = gr.Slider(minimum=100, maximum=50000, step=100, value=10000, label="Number of Samples for NCUT")
+                sampling_method_dropdown = gr.Dropdown(["fps", "random"], label="Sampling method", value="fps", elem_id="sampling_method")
+                knn_ncut = gr.Slider(minimum=1, maximum=100, step=1, value=10, label="KNN for NCUT")
+                embedding_method_dropdown = gr.Dropdown(["tsne_3d", "umap_3d", "umap_shpere", "tsne_2d", "umap_2d"], label="Coloring method", value="tsne_3d", elem_id="embedding_method")
+                num_sample_tsne_slider = gr.Slider(100, 10000, step=100, label="t-SNE/UMAP: num_sample", value=300, elem_id="num_sample_tsne", info="Nyström approximation")
+                knn_tsne_slider = gr.Slider(1, 100, step=1, label="t-SNE/UMAP: KNN", value=10, elem_id="knn_tsne", info="Nyström approximation")
+                perplexity_slider = gr.Slider(10, 1000, step=10, label="t-SNE: Perplexity", value=150, elem_id="perplexity")
+                n_neighbors_slider = gr.Slider(10, 1000, step=10, label="UMAP: n_neighbors", value=150, elem_id="n_neighbors")
+                min_dist_slider = gr.Slider(0.1, 1, step=0.1, label="UMAP: min_dist", value=0.1, elem_id="min_dist")
+            logging_str = gr.Textbox("", label="Logging Information", placeholder="Logging",)
+            
+    with gr.Row():
+        gr.Markdown("### Output Embedding")
+        output_image = gr.Plot(label="NCUT Output", min_width=1920)
+
+    def change_layer_slider(model_name):
+        layer_dict = TEXT_LAYER_DICT
+        if model_name in layer_dict:
+            value = layer_dict[model_name]
+            return (gr.Slider(1, value, step=1, label="Backbone: Layer index", value=value, elem_id="layer", visible=True),
+                    gr.Dropdown(["attn: attention output", "mlp: mlp output", "block: sum of residual"], label="Backbone: Layer type", value="block: sum of residual", elem_id="node_type", info="which feature to take from each layer?"))
+        else:
+            value = 12
+            return (gr.Dropdown(["attn: attention output", "mlp: mlp output", "block: sum of residual"], label="Backbone: Layer type", value="block: sum of residual", elem_id="node_type", info="which feature to take from each layer?"),
+                    gr.Slider(1, value, step=1, label="Backbone: Layer index", value=value, elem_id="layer", visible=True))
+    model_name.change(fn=change_layer_slider, inputs=model_name, outputs=[layer, node_type])
+
+    clear_button.click(lambda x: (None, None), outputs=[input_text, output_image])
+    submit_button.click(real_run, inputs=[
+        model_name, input_text, layer, node_type, num_eig, 
+        affinity_focal_gamma, num_sample_ncut, knn_ncut, 
+        embedding_method_dropdown, num_sample_tsne_slider, 
+        knn_tsne_slider, perplexity_slider, n_neighbors_slider, 
+        min_dist_slider, sampling_method_dropdown
+        ], 
+        outputs=[output_image, logging_str], 
+    )
+    
+
+
+if __name__ == "__main__":
+    with gr.Blocks() as demo:
+        make_demo()
+    demo.launch(share=True)

backbone_text.py

@@ -0,0 +1,239 @@
+# %%
+# 
+from typing import List, Union
+import torch
+import os
+from torch import nn
+from typing import Optional, Tuple
+
+from functools import partial
+
+MODEL_DICT = {}
+LAYER_DICT = {}
+
+class Llama(nn.Module):
+    def __init__(self, model_id="meta-llama/Meta-Llama-3.1-8B"):
+        super().__init__()
+        
+        import transformers
+        
+        access_token = os.getenv("HF_ACCESS_TOKEN")
+        if access_token is None:
+            raise ValueError("HF_ACCESS_TOKEN environment variable must be set")
+        
+        pipeline = transformers.pipeline(
+            "text-generation",
+            model=model_id,
+            model_kwargs={"torch_dtype": torch.bfloat16},
+            token=access_token,
+            device='cpu',
+        )
+        
+        tokenizer = pipeline.tokenizer
+        model = pipeline.model
+        
+        def new_forward(
+            self,
+            hidden_states: torch.Tensor,
+            attention_mask: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.LongTensor] = None,
+            past_key_value = None,
+            output_attentions: Optional[bool] = False,
+            use_cache: Optional[bool] = False,
+            cache_position: Optional[torch.LongTensor] = None,
+            position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # will become mandatory in v4.45
+            **kwargs,
+        ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+            residual = hidden_states
+
+            hidden_states = self.input_layernorm(hidden_states)
+
+            # Self Attention
+            hidden_states, self_attn_weights, present_key_value = self.self_attn(
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_value,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+            
+            self.attn_output = hidden_states.clone()
+            
+            hidden_states = residual + hidden_states
+
+            # Fully Connected
+            residual = hidden_states
+            hidden_states = self.post_attention_layernorm(hidden_states)
+            hidden_states = self.mlp(hidden_states)
+            
+            self.mlp_output = hidden_states.clone()
+            
+            hidden_states = residual + hidden_states
+            
+            self.block_output = hidden_states.clone()
+
+            outputs = (hidden_states,)
+
+            if output_attentions:
+                outputs += (self_attn_weights,)
+
+            if use_cache:
+                outputs += (present_key_value,)
+
+            return outputs
+        
+        # for layer in model.model.layers:
+            # setattr(layer.__class__, "forward", new_forward)
+            # setattr(layer.__class__, "__call__", new_forward)
+        setattr(model.model.layers[0].__class__, "forward", new_forward)
+        setattr(model.model.layers[0].__class__, "__call__", new_forward)
+        
+        self.model = model
+        self.tokenizer = tokenizer
+    
+    @torch.no_grad()
+    def forward(self, text: str):
+        encoded_input = self.tokenizer(text, return_tensors='pt')
+        device = next(self.model.parameters()).device
+        encoded_input = {k: v.to(device) for k, v in encoded_input.items()}
+        output = self.model(**encoded_input, output_hidden_states=True)
+    
+        attn_outputs, mlp_outputs, block_outputs = [], [], []
+        for i, blk in enumerate(self.model.model.layers):
+            attn_outputs.append(blk.attn_output)
+            mlp_outputs.append(blk.mlp_output)
+            block_outputs.append(blk.block_output)
+        
+        token_ids = encoded_input['input_ids']
+        token_texts = [self.tokenizer.decode([token_id]) for token_id in token_ids[0]]
+        
+        return {"attn": attn_outputs, "mlp": mlp_outputs, "block": block_outputs, "token_texts": token_texts}
+
+MODEL_DICT["meta-llama/Meta-Llama-3.1-8B"] = partial(Llama, model_id="meta-llama/Meta-Llama-3.1-8B")
+LAYER_DICT["meta-llama/Meta-Llama-3.1-8B"] = 32
+MODEL_DICT["meta-llama/Meta-Llama-3-8B"] = partial(Llama, model_id="meta-llama/Meta-Llama-3-8B")
+LAYER_DICT["meta-llama/Meta-Llama-3-8B"] = 32
+
+class GPT2(nn.Module):
+    def __init__(self):
+        super().__init__()
+        from transformers import GPT2Tokenizer, GPT2Model
+        tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
+        model = GPT2Model.from_pretrained('gpt2')
+        
+        def new_forward(
+            self,
+            hidden_states: Optional[Tuple[torch.FloatTensor]],
+            layer_past: Optional[Tuple[torch.Tensor]] = None,
+            attention_mask: Optional[torch.FloatTensor] = None,
+            head_mask: Optional[torch.FloatTensor] = None,
+            encoder_hidden_states: Optional[torch.Tensor] = None,
+            encoder_attention_mask: Optional[torch.FloatTensor] = None,
+            use_cache: Optional[bool] = False,
+            output_attentions: Optional[bool] = False,
+        ) -> Union[Tuple[torch.Tensor], Optional[Tuple[torch.Tensor, Tuple[torch.FloatTensor, ...]]]]:
+            residual = hidden_states
+            hidden_states = self.ln_1(hidden_states)
+            attn_outputs = self.attn(
+                hidden_states,
+                layer_past=layer_past,
+                attention_mask=attention_mask,
+                head_mask=head_mask,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+            )
+            attn_output = attn_outputs[0]  # output_attn: a, present, (attentions)
+            outputs = attn_outputs[1:]
+            # residual connection
+            self.attn_output = attn_output.clone()
+            hidden_states = attn_output + residual
+
+            if encoder_hidden_states is not None:
+                # add one self-attention block for cross-attention
+                if not hasattr(self, "crossattention"):
+                    raise ValueError(
+                        f"If `encoder_hidden_states` are passed, {self} has to be instantiated with "
+                        "cross-attention layers by setting `config.add_cross_attention=True`"
+                    )
+                residual = hidden_states
+                hidden_states = self.ln_cross_attn(hidden_states)
+                cross_attn_outputs = self.crossattention(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    head_mask=head_mask,
+                    encoder_hidden_states=encoder_hidden_states,
+                    encoder_attention_mask=encoder_attention_mask,
+                    output_attentions=output_attentions,
+                )
+                attn_output = cross_attn_outputs[0]
+                # residual connection
+                hidden_states = residual + attn_output
+                outputs = outputs + cross_attn_outputs[2:]  # add cross attentions if we output attention weights
+
+            residual = hidden_states
+            hidden_states = self.ln_2(hidden_states)
+            feed_forward_hidden_states = self.mlp(hidden_states)
+            # residual connection
+            self.mlp_output = feed_forward_hidden_states.clone()
+            hidden_states = residual + feed_forward_hidden_states
+
+            if use_cache:
+                outputs = (hidden_states,) + outputs
+            else:
+                outputs = (hidden_states,) + outputs[1:]
+
+            self.block_output = hidden_states.clone()
+            return outputs  # hidden_states, present, (attentions, cross_attentions)
+        
+        setattr(model.h[0].__class__, "forward", new_forward)
+        
+        self.model = model
+        self.tokenizer = tokenizer
+
+    @torch.no_grad()
+    def forward(self, text: str):
+        encoded_input = self.tokenizer(text, return_tensors='pt')
+        device = next(self.model.parameters()).device
+        encoded_input = {k: v.to(device) for k, v in encoded_input.items()}
+        output = self.model(**encoded_input, output_hidden_states=True)
+    
+        attn_outputs, mlp_outputs, block_outputs = [], [], []
+        for i, blk in enumerate(self.model.h):
+            attn_outputs.append(blk.attn_output)
+            mlp_outputs.append(blk.mlp_output)
+            block_outputs.append(blk.block_output)
+            
+        token_ids = encoded_input['input_ids']
+        token_texts = [self.tokenizer.decode([token_id]) for token_id in token_ids[0]]
+        
+        return {"attn": attn_outputs, "mlp": mlp_outputs, "block": block_outputs, "token_texts": token_texts}
+    
+MODEL_DICT["gpt2"] = GPT2
+LAYER_DICT["gpt2"] = 12
+
+
+def download_all_models():
+    for model_name in MODEL_DICT:
+        print(f"Downloading {model_name}")
+        try:
+            model = MODEL_DICT[model_name]()
+        except Exception as e:
+            print(f"Error downloading {model_name}: {e}")
+            continue
+
+
+if __name__ == '__main__':
+
+    model = MODEL_DICT["meta-llama/Meta-Llama-3-8B"]()
+    # model = MODEL_DICT["gpt2"]()
+    text = """
+    1. The majestic giraffe, with its towering height and distinctive long neck, roams the savannas of Africa. These gentle giants use their elongated tongues to pluck leaves from the tallest trees, making them well-adapted to their environment. Their unique coat patterns, much like human fingerprints, are unique to each individual.
+    """
+    model = model.cuda()
+    output = model(text)
+    print(output["block"][1].shape)
+    print(output["token_texts"])