fix(ci): fix tests and model compatibility for CI/CD

- Limit transformers dependency to <5.0.0 to prevent breaking changes in model outputs.
- Fix LlamaRotaryEmbedding instantiation in LlamaAdapter by passing the correct dim.
- Update test imports from src.lema to lema.

Co-authored-by: Pomilon <220483426+Pomilon@users.noreply.github.com>

Author: google-labs-jules[bot]

pyproject.toml

@@ -27,7 +27,7 @@ classifiers = [
 ]
 dependencies = [
     "torch>=2.0.0",
-    "transformers>=4.30.0",
+    "transformers>=4.30.0,<5.0.0",
     "safetensors>=0.3.0",
     "accelerate>=0.20.0",
     "peft>=0.4.0",

requirements.txt

@@ -1,5 +1,5 @@
 torch>=2.0.0
-transformers>=4.30.0
+transformers>=4.30.0,<5.0.0
 safetensors>=0.3.0
 accelerate>=0.20.0
 peft>=0.4.0

src/lema/models/llama.py

@@ -11,7 +11,7 @@ def __init__(self, config: Dict[str, Any]):
         self.hf_config = LlamaConfig(**config)
         if getattr(self.hf_config, "_attn_implementation", None) is None:
             self.hf_config._attn_implementation = config.get("attn_implementation", "eager")
-        self.rotary_emb = LlamaRotaryEmbedding(self.hf_config)
+        self.rotary_emb = LlamaRotaryEmbedding(self.hf_config.hidden_size // self.hf_config.num_attention_heads, max_position_embeddings=self.hf_config.max_position_embeddings)
         self.layer_pool: List[nn.Module] = []
         self.param_mappings: Dict[int, List[tuple]] = {}
         self._max_pool_size = 8

tests/test_core_components.py

@@ -1,9 +1,9 @@
 import pytest
 import torch
 import torch.nn as nn
-from src.lema.core.gbi import GlobalBinaryIndex
-from src.lema.core.lora import LoRAManager, LoRAWrapper
-from src.lema.core.memory import TripleBufferManager
+from lema.core.gbi import GlobalBinaryIndex
+from lema.core.lora import LoRAManager, LoRAWrapper
+from lema.core.memory import TripleBufferManager
 
 # Mocking
 class MockAdapter:

tests/test_gradient_equivalence.py

@@ -2,10 +2,10 @@
 import torch.nn as nn
 from transformers import GPT2Config, GPT2LMHeadModel
 from safetensors.torch import save_file
-from src.lema.core.gbi import GlobalBinaryIndex
-from src.lema.models.gpt2 import GPT2Adapter
-from src.lema.engine.trainer import LemaTrainer
-from src.lema.config import LemaConfig, MemoryStrategy
+from lema.core.gbi import GlobalBinaryIndex
+from lema.models.gpt2 import GPT2Adapter
+from lema.engine.trainer import LemaTrainer
+from lema.config import LemaConfig, MemoryStrategy
 import os
 import pytest
 

tests/test_llama_forward.py

@@ -1,6 +1,6 @@
 import torch
 import torch.nn as nn
-from src.lema.models.llama import LlamaAdapter
+from lema.models.llama import LlamaAdapter
 from transformers import LlamaConfig
 
 def test_llama_adapter_forward():