Update gradio_app.py

gradio_app.py

@@ -41,7 +41,12 @@ DEVICE = os.environ.get("DEVICE", "cpu")  # cuda:0
 if DEVICE != "cpu" and not torch.cuda.is_available():
     DEVICE = "cpu"
 logger.info(f"DEVICE {DEVICE}")
-DTYPE = torch.float32 if DEVICE == "cpu" else torch.float16
+DTYPE = getattr(
+    torch,
+    os.environ.get("DTYPE", ""),
+    torch.float32 if DEVICE == "cpu" else torch.float16
+)
+LOW_CPU_MEM = bool(os.environ.get("LOW_CPU_MEM", False if DEVICE == "cpu" else True))
 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))
@@ -147,7 +152,7 @@ class TextGeneration:
             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
+            torch_dtype=DTYPE, low_cpu_mem_usage=LOW_CPU_MEM,
         ).to(device=DEVICE, non_blocking=False)
         _ = self.model.eval()
         device_number = -1 if DEVICE == "cpu" else int(DEVICE.split(":")[-1])