Update app.py

app.py

@@ -15,22 +15,6 @@ print(accelerate.__version__)    # Should be >= 0.12.0
 num_gpus = torch.cuda.device_count()
 print(f"Number of available GPUs: {num_gpus}")
 
-# List details for each GPU
-for i in range(num_gpus):
-    print(f"GPU {i}: {torch.cuda.get_device_name(i)}")
-    print(f"  Total Memory: {torch.cuda.get_device_properties(i).total_memory / 1e9:.2f} GB")
-    print(f"  CUDA Capability: {torch.cuda.get_device_properties(i).major}.{torch.cuda.get_device_properties(i).minor}")
-for i in range(num_gpus):
-    print(f"GPU {i}: {torch.cuda.get_device_name(i)}")
-    print(f"  Allocated Memory: {torch.cuda.memory_allocated(i) / 1e9:.2f} GB")
-    print(f"  Cached Memory: {torch.cuda.memory_reserved(i) / 1e9:.2f} GB")
-    print(f"  Free Memory: {torch.cuda.get_device_properties(i).total_memory / 1e9 - torch.cuda.memory_reserved(i) / 1e9:.2f} GB")
-
-quantization_config = BitsAndBytesConfig(
-    load_in_8bit=True,  # Enable 8-bit quantization
-    llm_int8_enable_fp32_cpu_offload=True  # Enable FP32 CPU offloading
-)
-
 """
 For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
 """
@@ -71,6 +55,24 @@ PROMPT_DICT = {
 }
 model = None
 tokenizer = None
+
+quantization_config = BitsAndBytesConfig(
+    load_in_8bit=True,  # Enable 8-bit quantization
+    llm_int8_enable_fp32_cpu_offload=True  # Enable FP32 CPU offloading
+)
+
+def print_resources():
+    # List details for each GPU
+    for i in range(num_gpus):
+        print(f"GPU {i}: {torch.cuda.get_device_name(i)}")
+        print(f"  Total Memory: {torch.cuda.get_device_properties(i).total_memory / 1e9:.2f} GB")
+        print(f"  CUDA Capability: {torch.cuda.get_device_properties(i).major}.{torch.cuda.get_device_properties(i).minor}")
+    for i in range(num_gpus):
+        print(f"GPU {i}: {torch.cuda.get_device_name(i)}")
+        print(f"  Allocated Memory: {torch.cuda.memory_allocated(i) / 1e9:.2f} GB")
+        print(f"  Cached Memory: {torch.cuda.memory_reserved(i) / 1e9:.2f} GB")
+        print(f"  Free Memory: {torch.cuda.get_device_properties(i).total_memory / 1e9 - torch.cuda.memory_reserved(i) / 1e9:.2f} GB")
+        
 def generate_prompt(instruction, input=None):
     if input:
         return PROMPT_DICT["prompt_input"].format(instruction=instruction,input=input)
@@ -118,6 +120,7 @@ def loadModel():
         #        )
         tokenizer =  AutoTokenizer.from_pretrained(base_model,use_fast=False,cache_dir=cache_dir)
         tokenizer.pad_token = tokenizer.unk_token
+    print_resources()
     return model, tokenizer
 
 model, tokenizer = loadModel()
@@ -130,11 +133,13 @@ def respond(
     max_tokens,
     temperature,
     top_p,
-):    
+):
     ins_f = generate_prompt(message,None)
     inputs  =  tokenizer(ins_f, return_tensors="pt")
+    print_resources()
     input_ids = inputs["input_ids"].cuda()
     max_new_tokens = 512
+    print_resources()
     generation_config = GenerationConfig(
             temperature=0.1,
             top_p=0.75,