Update app.py

app.py

@@ -18,11 +18,29 @@ subprocess.run(
    
     shell=True,
 )
+
+subprocess.run(
+    "pip install pynvml gpustat",
+   
+    shell=True,
+)
 def print_s1ystem():
     ram_info = psutil.virtual_memory()
     print(f"Total RAM: {ram_info.total / (1024.0 ** 3)} GB")
     print(f"Available RAM: {ram_info.available / (1024.0 ** 3)} GB")
 
+import psutil
+import platform
+import gpustat
+from datetime import datetime
+
+def get_size(bytes, suffix="B"):
+    factor = 1024
+    for unit in ["", "K", "M", "G", "T", "P"]:
+        if bytes < factor:
+            return f"{bytes:.2f}{unit}{suffix}"
+        bytes /= factor
+
 subprocess.run(
     "pip install flash-attn --no-build-isolation",
     env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
@@ -84,6 +102,18 @@ def respond(
     top_p,
 ):
     print_s1ystem()
+    uname = platform.uname()
+    print(f"System: {uname.system}")
+    print(f"Node Name: {uname.node}")
+    print(f"Release: {uname.release}")
+    print(f"Version: {uname.version}")
+    print(f"Machine: {uname.machine}")
+    print(f"Processor: {uname.processor}")
+
+# GPU Information
+gpu_stats = gpustat.GPUStatCollection.new_query()
+for gpu in gpu_stats:
+    print(f"GPU: {gpu.name}  Mem Free: {get_size(gpu.memory_free)}  Mem Used: {get_size(gpu.memory_used)}  Mem Total: {get_size(gpu.memory_total)}")
     messages = [{"role": "user", "content": "Hello, how are you?"}]
     input_ids = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to('cuda')
 ## <BOS_TOKEN><|START_OF_TURN_TOKEN|><|USER_TOKEN|>Hello, how are you?<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>