Fixed tests

predibase · Nov 1, 2024 · 295829f · 295829f
1 parent 35c7de2
commit 295829f
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/server/tests/adapters/test_medusa.py b/server/tests/adapters/test_medusa.py
@@ -30,6 +30,7 @@ def test_batched_medusa_weights(default_causal_lm: CausalLM):
 
     meta = AdapterBatchMetadata(
         adapter_indices=torch.tensor([0, 0, 1, 1, 0, 0, 1, 1], dtype=torch.int64),
+        adapter_list=[0, 1, 0, 1],
         adapter_set={0, 1},
         adapter_segments=torch.tensor([0, 2, 4, 6, 8], dtype=torch.int64),
         segment_indices=[0, 1, 0, 1],

diff --git a/server/tests/utils/test_lora.py b/server/tests/utils/test_lora.py
@@ -74,6 +74,7 @@ def test_batched_lora_weights(lora_ranks: List[int]):
 
     meta = AdapterBatchMetadata(
         adapter_indices=torch.tensor([0, 0, 1, 1, 0, 0, 1, 1], dtype=torch.int64),
+        adapter_list=[0, 1, 0, 1],
         adapter_set={0, 1},
         adapter_segments=torch.tensor([0, 2, 4, 6, 8], dtype=torch.int64),
         segment_indices=[0, 1, 0, 1],
@@ -149,6 +150,7 @@ def test_batched_lora_weights_decode(
 
     meta = AdapterBatchMetadata(
         adapter_indices=torch.tensor(adapter_indices, dtype=torch.int64),
+        adapter_list=adapter_indices,
         adapter_set=set(adapter_indices),
         adapter_segments=torch.tensor(segments, dtype=torch.int64),
         segment_indices=segment_indices,
@@ -193,7 +195,8 @@ def test_batched_lora_weights_no_segments():
 
     meta = AdapterBatchMetadata(
         adapter_indices=torch.tensor([0, 0, 0, 0], dtype=torch.int64),
-        adapter_set={0, 1},
+        adapter_list=[0],
+        adapter_set={0},
         adapter_segments=torch.tensor([0, 4], dtype=torch.int64),
         segment_indices=[0],
     )