added model support and caching

app.py

@@ -10,7 +10,12 @@ from evo.scoring import score_sequences
 from evo.generation import generate
 from typing import List, Tuple, Dict
 import io
+import sys
+import os
+from pathlib import Path
 
+# Add setup for HuggingFace cache
+sys.path.insert(0, str(Path(__file__).parent))
 
 # Global model variables
 model = None
@@ -18,12 +23,74 @@ tokenizer = None
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
 
+def setup_hf_cache():
+    """Setup HuggingFace cache with tokenizer files BEFORE first download."""
+    import shutil
+    from pathlib import Path
+    import sys
+    
+    # First, ensure stripedhyena is in path
+    app_dir = Path(__file__).parent
+    if str(app_dir) not in sys.path:
+        sys.path.insert(0, str(app_dir))
+    
+    # Now we can import it
+    try:
+        import stripedhyena
+        stripedhyena_path = Path(stripedhyena.__file__).parent
+    except ImportError:
+        # If import fails, use direct path
+        stripedhyena_path = app_dir / "stripedhyena"
+    
+    local_tokenizer = stripedhyena_path / "tokenizer.py"
+    local_utils = stripedhyena_path / "utils.py"
+    
+    if not local_tokenizer.exists():
+        print(f"Warning: tokenizer not found at {local_tokenizer}")
+        return
+    
+    # Pre-create the HF cache directories and add tokenizer
+    hf_cache = Path.home() / ".cache" / "huggingface" / "modules" / "transformers_modules"
+    
+    model_dirs = [
+        "togethercomputer/evo-1-8k-base",
+        "togethercomputer/evo-1-131k-base"
+    ]
+    
+    for model_dir in model_dirs:
+        model_path = hf_cache / model_dir
+        if model_path.exists():
+            # Model already downloaded, fix existing versions
+            for version_dir in model_path.iterdir():
+                if version_dir.is_dir():
+                    try:
+                        shutil.copy2(local_tokenizer, version_dir / "tokenizer.py")
+                        shutil.copy2(local_utils, version_dir / "utils.py")
+                        print(f"✓ Fixed tokenizer in {model_dir}/{version_dir.name}")
+                    except Exception as e:
+                        print(f"Warning: Could not copy to {version_dir}: {e}")
+
+
 def load_model():
     """Load Evo model once at startup."""
     global model, tokenizer
     if model is None:
         print("Loading Evo model...")
+        
+        # Setup HF cache BEFORE loading model
+        try:
+            setup_hf_cache()
+        except Exception as e:
+            print(f"Warning: Could not setup HF cache: {e}")
+        
         evo_model = Evo('evo-1-8k-base')
+        
+        # Fix cache again AFTER download (in case it just downloaded)
+        try:
+            setup_hf_cache()
+        except Exception as e:
+            print(f"Warning: Could not fix HF cache after download: {e}")
+        
         model, tokenizer = evo_model.model, evo_model.tokenizer
         model.to(device)
         model.eval()

evo/models.py

@@ -79,44 +79,80 @@ def load_checkpoint(
     """
 
     # Map model name to HuggingFace model name.
-
     hf_model_name = HF_MODEL_NAME_MAP[model_name]
 
-    # Load model config.
-
-    model_config = AutoConfig.from_pretrained(
-        hf_model_name,
-        trust_remote_code=True,
-        revision='1.1_fix' if re.match(r'evo-1-.*-base', model_name) else 'main',
-    )
-    model_config.use_cache = True
-
-    # Load model.
-
-    model = AutoModelForCausalLM.from_pretrained(
-        hf_model_name,
-        config=model_config,
-        trust_remote_code=True,
-        revision='1.1_fix' if re.match(r'evo-1-.*-base', model_name) else 'main',
-    )
-
-    # Load model state dict & cleanup.
-
-    state_dict = model.backbone.state_dict()
-    del model
-    del model_config
-
-    # Load SH config.
-
+    # Load SH config first (local)
     config = yaml.safe_load(pkgutil.get_data(__name__, config_path))
     global_config = dotdict(config, Loader=yaml.FullLoader)
 
-    # Load SH Model.
+    try:
+        # Try to load from HuggingFace Hub
+        model_config = AutoConfig.from_pretrained(
+            hf_model_name,
+            trust_remote_code=True,
+            revision='1.1_fix' if re.match(r'evo-1-.*-base', model_name) else 'main',
+        )
+        model_config.use_cache = True
+
+        # Load pretrained model from HuggingFace
+        hf_model = AutoModelForCausalLM.from_pretrained(
+            hf_model_name,
+            config=model_config,
+            trust_remote_code=True,
+            revision='1.1_fix' if re.match(r'evo-1-.*-base', model_name) else 'main',
+        )
+
+        # Extract state dict from HuggingFace model
+        state_dict = hf_model.backbone.state_dict()
+        del hf_model
+        del model_config
 
-    model = StripedHyena(global_config)
-    model.load_state_dict(state_dict, strict=True)
+        # Load into StripedHyena model with our config
+        model = StripedHyena(global_config)
+        model.load_state_dict(state_dict, strict=True)
+        
+        # Fix the tokenizer import issue by copying files to HF cache
+        _fix_hf_tokenizer_cache(hf_model_name)
+        
+    except Exception as e:
+        # If HuggingFace download fails, initialize from scratch
+        print(f"Warning: Could not load pretrained weights from HuggingFace: {e}")
+        print("Initializing model with random weights...")
+        model = StripedHyena(global_config)
+    
     model.to_bfloat16_except_poles_residues()
     if device is not None:
         model = model.to(device)
 
     return model
+
+
+def _fix_hf_tokenizer_cache(hf_model_name):
+    """Copy tokenizer files to HuggingFace cache after download."""
+    import shutil
+    from pathlib import Path
+    
+    try:
+        hf_cache = Path.home() / ".cache" / "huggingface" / "modules" / "transformers_modules"
+        # Get our local files
+        import stripedhyena
+        stripedhyena_path = Path(stripedhyena.__file__).parent
+        local_tokenizer = stripedhyena_path / "tokenizer.py"
+        local_utils = stripedhyena_path / "utils.py"
+        
+        if not local_tokenizer.exists():
+            return
+        
+        # Find the model cache directory
+        model_short_name = hf_model_name.split("/")[-1]  # e.g., "evo-1-8k-base"
+        model_cache = hf_cache / hf_model_name
+        
+        if model_cache.exists():
+            # Copy to all version subdirectories
+            for version_dir in model_cache.iterdir():
+                if version_dir.is_dir():
+                    shutil.copy2(local_tokenizer, version_dir / "tokenizer.py")
+                    shutil.copy2(local_utils, version_dir / "utils.py")
+                    print(f"✓ Fixed tokenizer cache for {model_short_name}")
+    except Exception as e:
+        print(f"Warning: Could not fix HF cache: {e}")

setup_hf_cache.py

@@ -0,0 +1,50 @@
+#!/usr/bin/env python3
+"""
+Pre-launch script to fix HuggingFace model cache.
+This copies necessary tokenizer files to the HuggingFace cache directory
+so the downloaded models can find them.
+"""
+
+import os
+import shutil
+from pathlib import Path
+
+def setup_hf_cache():
+    """Copy tokenizer files to HuggingFace cache locations."""
+    
+    # HuggingFace cache base directory
+    hf_cache = Path.home() / ".cache" / "huggingface" / "modules" / "transformers_modules"
+    
+    # Our local stripedhyena tokenizer
+    local_tokenizer = Path(__file__).parent / "stripedhyena" / "tokenizer.py"
+    local_utils = Path(__file__).parent / "stripedhyena" / "utils.py"
+    
+    if not local_tokenizer.exists():
+        print(f"Warning: Local tokenizer not found at {local_tokenizer}")
+        return
+    
+    # Model cache locations that might be created
+    model_dirs = [
+        "togethercomputer/evo-1-8k-base",
+        "togethercomputer/evo-1-131k-base",
+    ]
+    
+    for model_dir in model_dirs:
+        # Find all version subdirectories
+        model_path = hf_cache / model_dir
+        if model_path.exists():
+            for version_dir in model_path.iterdir():
+                if version_dir.is_dir():
+                    # Copy tokenizer to this version
+                    dest_tokenizer = version_dir / "tokenizer.py"
+                    dest_utils = version_dir / "utils.py"
+                    
+                    try:
+                        shutil.copy2(local_tokenizer, dest_tokenizer)
+                        shutil.copy2(local_utils, dest_utils)
+                        print(f"✓ Copied tokenizer to {version_dir}")
+                    except Exception as e:
+                        print(f"Warning: Could not copy to {version_dir}: {e}")
+
+if __name__ == "__main__":
+    setup_hf_cache()