Upload 50 files

.gitattributes

@@ -1,2 +1,3 @@
 static/favicon.ico filter=lfs diff=lfs merge=lfs -text
 static/images/architecture.png filter=lfs diff=lfs merge=lfs -text
+attention_app/server/__pycache__/main.cpython-311.pyc filter=lfs diff=lfs merge=lfs -text

attention_app/models.py

@@ -0,0 +1,108 @@
+import logging
+import warnings
+import torch
+from transformers import BertTokenizer, BertModel, BertForMaskedLM
+from transformers.utils import logging as transformers_logging
+from transformers import GPT2Tokenizer, GPT2Model, GPT2LMHeadModel
+import gc
+
+# Suppress warnings
+warnings.filterwarnings(
+    "ignore",
+    message="`resume_download` is deprecated",
+    category=FutureWarning,
+    module="huggingface_hub.file_download",
+)
+transformers_logging.set_verbosity_error()
+logging.getLogger("transformers").setLevel(logging.ERROR)
+
+
+class ModelManager:
+    """Manages loading and caching of BERT models."""
+    
+    _instances = {}
+
+    @classmethod
+    def get_model(cls, model_name: str):
+        """
+        Returns (tokenizer, encoder_model, mlm_model) for the specified model_name.
+        Loads from cache if available, otherwise loads from HuggingFace.
+        """
+        # Check if model is already loaded
+        if model_name in cls._instances:
+            return cls._instances[model_name]
+
+        # Clear existing cache to free memory
+        if cls._instances:
+            print(f"Unloading previous models to free memory...")
+            cls._instances.clear()
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
+        print(f"Loading model: {model_name}...")
+        
+        try:
+
+            is_gpt2 = "gpt2" in model_name
+            
+            if is_gpt2:
+                tokenizer = GPT2Tokenizer.from_pretrained(model_name)
+                if tokenizer.pad_token is None:
+                    tokenizer.pad_token = tokenizer.eos_token
+                    
+                encoder = GPT2Model.from_pretrained(
+                    model_name,
+                    output_attentions=True,
+                    output_hidden_states=True,
+                )
+                encoder.eval()
+                
+                try:
+                    mlm = GPT2LMHeadModel.from_pretrained(
+                        model_name,
+                        output_attentions=False,
+                        output_hidden_states=False,
+                    )
+                    mlm.eval()
+                except Exception as e:
+                    print(f"Warning: Could not load LM head for {model_name}: {e}")
+                    mlm = None
+            else:
+                tokenizer = BertTokenizer.from_pretrained(model_name)
+                
+                encoder = BertModel.from_pretrained(
+                    model_name,
+                    output_attentions=True,
+                    output_hidden_states=True,
+                )
+                encoder.eval()
+                
+                try:
+                    mlm = BertForMaskedLM.from_pretrained(
+                        model_name,
+                        output_attentions=False,
+                        output_hidden_states=False,
+                    )
+                    mlm.eval()
+                except Exception as e:
+                    print(f"Warning: Could not load MLM head for {model_name}: {e}")
+                    mlm = None
+
+            # Move to GPU if available
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            encoder.to(device)
+            if mlm:
+                mlm.to(device)
+
+            cls._instances[model_name] = (tokenizer, encoder, mlm)
+            return tokenizer, encoder, mlm
+            
+        except Exception as e:
+            raise RuntimeError(f"Failed to load model {model_name}: {str(e)}")
+
+    @staticmethod
+    def get_device():
+        return "cuda" if torch.cuda.is_available() else "cpu"
+
+__all__ = ["ModelManager"]

attention_app/server/__pycache__/main.cpython-311.pyc

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9fbca4c333ba02d6b4139d59210196392c9f834097a8aaca78c574bd7fe9e564
+size 101097

attention_app/ui/__init__.py

@@ -0,0 +1,5 @@
+"""UI package for Attention Atlas."""
+
+from .layouts import app_ui
+
+__all__ = ["app_ui"]