From b5790224a536d1a2a9fdae655c0f1f4c516e0b21 Mon Sep 17 00:00:00 2001
From: Nicolas Patry <patry.nicolas@protonmail.com>
Date: Thu, 23 Nov 2023 12:55:18 +0100
Subject: [PATCH] Apply suggestions from code review

---
 server/text_generation_server/cli.py | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/server/text_generation_server/cli.py b/server/text_generation_server/cli.py
index c8f40c46..b741a84c 100644
--- a/server/text_generation_server/cli.py
+++ b/server/text_generation_server/cli.py
@@ -150,15 +150,16 @@ def download_weights(
             if not extension == ".safetensors" or not auto_convert:
                 raise e
 
-    # Try to load as a local PEFT model
-    try:
-        utils.download_and_unload_peft(
-            model_id, revision, trust_remote_code=trust_remote_code
-        )
-        utils.weight_files(model_id, revision, extension)
-        return
-    except (utils.LocalEntryNotFoundError, utils.EntryNotFoundError):
-        pass
+    else:
+        # Try to load as a local PEFT model
+        try:
+            utils.download_and_unload_peft(
+                model_id, revision, trust_remote_code=trust_remote_code
+            )
+            utils.weight_files(model_id, revision, extension)
+            return
+        except (utils.LocalEntryNotFoundError, utils.EntryNotFoundError):
+            pass
 
     # Try to see if there are local pytorch weights
     try: