From bfec09ecc2ccc9e22eeb359d7d4e448410455e45 Mon Sep 17 00:00:00 2001
From: Nicolas Patry <patry.nicolas@protonmail.com>
Date: Mon, 26 Feb 2024 16:31:40 +0000
Subject: [PATCH] Fixing revision for the medusa test.

---
 server/text_generation_server/models/__init__.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/server/text_generation_server/models/__init__.py b/server/text_generation_server/models/__init__.py
index 9c0a5eaa..3208275c 100644
--- a/server/text_generation_server/models/__init__.py
+++ b/server/text_generation_server/models/__init__.py
@@ -124,6 +124,7 @@ def get_model(
     use_medusa = None
     if "medusa_num_heads" in config_dict:
         medusa_model_id = model_id
+        medusa_revision = revision
         model_id = config_dict["base_model_name_or_path"]
         revision = "main"
         speculate_medusa = config_dict["medusa_num_heads"]
@@ -143,11 +144,11 @@ def get_model(
         is_local = Path(medusa_model_id).exists()
         if not is_local:
             medusa_config = hf_hub_download(
-                medusa_model_id, revision=revision, filename="config.json"
+                medusa_model_id, revision=medusa_revision, filename="config.json"
             )
             hf_hub_download(
                 medusa_model_id,
-                revision=revision,
+                revision=medusa_revision,
                 filename="medusa_lm_head.safetensors",
             )
             use_medusa = Path(medusa_config).parent