Update gradio_app.py

gradio_app.py

@@ -43,6 +43,7 @@ if DEVICE != "cpu" and not torch.cuda.is_available():
 logger.info(f"DEVICE {DEVICE}")
 DTYPE = torch.float32 if DEVICE == "cpu" else torch.float16
 MODEL_NAME = os.environ.get("MODEL_NAME", "bertin-project/bertin-gpt-j-6B")
+MODEL_REVISION = os.environ.get("MODEL_REVISION", "main")
 MAX_LENGTH = int(os.environ.get("MAX_LENGTH", 1024))
 HEADER_INFO = """
 # BERTIN GPT-J-6B
@@ -140,10 +141,11 @@ class TextGeneration:
     def load(self):
         logger.info("Loading model...")
         self.tokenizer = AutoTokenizer.from_pretrained(
-            self.model_name_or_path, use_auth_token=HF_AUTH_TOKEN if HF_AUTH_TOKEN else None,
+            self.model_name_or_path, revision=MODEL_REVISION, use_auth_token=HF_AUTH_TOKEN if HF_AUTH_TOKEN else None,
         )
         self.model = AutoModelForCausalLM.from_pretrained(
-            self.model_name_or_path, use_auth_token=HF_AUTH_TOKEN if HF_AUTH_TOKEN else None,
+            self.model_name_or_path, revision=MODEL_REVISION,
+            use_auth_token=HF_AUTH_TOKEN if HF_AUTH_TOKEN else None,
             pad_token_id=self.tokenizer.eos_token_id, eos_token_id=self.tokenizer.eos_token_id,
             torch_dtype=DTYPE, low_cpu_mem_usage=False if DEVICE == "cpu" else True
         ).to(device=DEVICE, non_blocking=False)