add more system info and detect HF spaces deployment

app.py

@@ -10,23 +10,34 @@ import logging
 from io import StringIO
 import time
 import asyncio
-# Create a string stream to capture log messages
-log_stream = StringIO()
-
-# Configure logging to use the string stream
-logging.basicConfig(stream=log_stream, level=logging.DEBUG, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
-logger = logging.getLogger(__name__)
-log_contents = log_stream.getvalue()
-print(log_contents)
-logger.debug('This is a debug message')
+import psutil
+import platform
+import os
 
 hf_user = None
-hfApi = HfApi()
 try:
+    hfApi = HfApi()
     hf_user = hfApi.whoami()["name"]
 except Exception as e:
     hf_user = "not logged in"
 
+def get_human_readable_size(size, decimal_places=2):
+    for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
+        if size < 1024.0:
+            break
+        size /= 1024.0
+    return f"{size:.{decimal_places}f} {unit}"
+
+
+# get cpu stats
+disk_stats = psutil.disk_usage('.')
+print(get_human_readable_size(disk_stats.total))
+cpu_info = platform.processor()
+print(cpu_info)
+os_info = platform.platform()
+print(os_info)
+
+memory = psutil.virtual_memory()
 
 # Dropdown options
 model_options = [
@@ -55,6 +66,18 @@ gpu_stats = torch.cuda.get_device_properties(0)
 start_gpu_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)
 max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3)
 
+running_on_hf = False
+if os.getenv("SYSTEM", None) == "spaces":
+    running_on_hf = True
+
+system_info = f"""\
+- **System:** {os_info}
+- **CPU:** {cpu_info} **Memory:** {get_human_readable_size(memory.free)} free of {get_human_readable_size(memory.total)}
+- **GPU:** {gpu_stats.name} ({max_memory} GB)
+- **Disk:** {get_human_readable_size(disk_stats.free)} free of {get_human_readable_size(disk_stats.total)}
+- **Hugging Face:** {running_on_hf}
+"""
+
 model=None
 tokenizer = None
 dataset = None
@@ -100,8 +123,6 @@ def load_model(initial_model_name, load_in_4bit, max_sequence_length):
         load_in_4bit = load_in_4bit,
         # token = "hf_...", # use one if using gated models like meta-llama/Llama-2-7b-hf
     )
-    log_contents = log_stream.getvalue()
-    print(log_contents)
     return f"Model {initial_model_name} loaded, using {max_sequence_length} as max sequence length.", gr.update(visible=True, interactive=True), gr.update(interactive=True),gr.update(interactive=False), gr.update(interactive=False), gr.update(interactive=False)
 
 def load_data(dataset_name, data_template_style, data_template):
@@ -162,17 +183,22 @@ def save_model(model_name, hub_model_name, hub_token, gguf_16bit, gguf_8bit, ggu
     return "Model saved", gr.update(visible=True, interactive=True)
 
 def username(profile: gr.OAuthProfile | None):
-    return profile["name"] if profile else "not logged in"
+    hf_user = profile["name"] if profile else "not logged in"
+    return hf_user
 
 # Create the Gradio interface
 with gr.Blocks(title="Unsloth fine-tuning") as demo:
-    gr.LoginButton()
-    logged_user = gr.Markdown(f"**User:** {hf_user}")
-    demo.load(username, inputs=None, outputs=logged_user)
-    with gr.Column():
-        gr.Image("unsloth.png", width="300px", interactive=False, show_download_button=False, show_label=False)
-    with gr.Column():
-        gr.Markdown(f"**User:** {hf_user}\n\n**GPU Information:** {gpu_stats.name} ({max_memory} GB)\n\n[Unsloth Docs](http://docs.unsloth.com/)\n\n[Unsloth GitHub](https://github.com/unslothai/unsloth)") 
+    if (running_on_hf):
+        gr.LoginButton()
+    # logged_user = gr.Markdown(f"**User:** {hf_user}")
+    #demo.load(username, inputs=None, outputs=logged_user)
+    with gr.Row():
+        with gr.Column(scale=0.5):
+            gr.Image("unsloth.png", width="300px", interactive=False, show_download_button=False, show_label=False, show_share_button=False)
+        with gr.Column(min_width="550px", scale=1):
+            gr.Markdown(system_info) 
+        with gr.Column(min_width="250px", scale=0.3):
+            gr.Markdown(f"**Links:**\n\n* [Unsloth Hub](https://huggingface.co/unsloth)\n\n* [Unsloth Docs](http://docs.unsloth.com/)\n\n* [Unsloth GitHub](https://github.com/unslothai/unsloth)")
     with gr.Tab("Base Model Parameters"):
 
         with gr.Row():
@@ -232,12 +258,12 @@ with gr.Blocks(title="Unsloth fine-tuning") as demo:
                     log_to_tensorboard = gr.Checkbox(label="Log to Tensorboard", value=True, interactive=True)
 
                 with gr.Row():
-                    optim = gr.Dropdown(choices=["adamw_8bit", "adamw", "sgd"], label="Optimizer", value="adamw_8bit")
+                #     optim = gr.Dropdown(choices=["adamw_8bit", "adamw", "sgd"], label="Optimizer", value="adamw_8bit")
                     learning_rate = gr.Number(label="Learning Rate", value=2e-4, interactive=True)
 
-                with gr.Row():
+                # with gr.Row():
                     weight_decay = gr.Number(label="Weight Decay", value=0.01, interactive=True)
-                    lr_scheduler_type = gr.Dropdown(choices=["linear", "cosine", "constant"], label="LR Scheduler Type", value="linear")
+                    # lr_scheduler_type = gr.Dropdown(choices=["linear", "cosine", "constant"], label="LR Scheduler Type", value="linear")
         gr.Markdown("---")
 
         with gr.Row():
@@ -249,9 +275,9 @@ with gr.Blocks(title="Unsloth fine-tuning") as demo:
         train_btn = gr.Button("Train", visible=True)
 
         def train_model(model_name: str, lora_r: int, lora_alpha: int, lora_dropout: float, per_device_train_batch_size: int, warmup_steps: int, max_steps: int,
-                gradient_accumulation_steps: int, logging_steps: int, log_to_tensorboard: bool, optim, learning_rate, weight_decay, lr_scheduler_type, seed: int, output_dir, progress= gr.Progress()):
+                gradient_accumulation_steps: int, logging_steps: int, log_to_tensorboard: bool, learning_rate, weight_decay, seed: int, output_dir, progress= gr.Progress()):
             global model, tokenizer
-            print(f"$$$ Training model {model_name} with {lora_r} R, {lora_alpha} alpha, {lora_dropout} dropout, {per_device_train_batch_size} per device train batch size, {warmup_steps} warmup steps, {max_steps} max steps, {gradient_accumulation_steps} gradient accumulation steps, {logging_steps} logging steps, {log_to_tensorboard} log to tensorboard, {optim} optimizer, {learning_rate} learning rate, {weight_decay} weight decay, {lr_scheduler_type} lr scheduler type, {seed} seed, {output_dir} output dir")
+            print(f"$$$ Training model {model_name} with {lora_r} R, {lora_alpha} alpha, {lora_dropout} dropout, {per_device_train_batch_size} per device train batch size, {warmup_steps} warmup steps, {max_steps} max steps, {gradient_accumulation_steps} gradient accumulation steps, {logging_steps} logging steps, {log_to_tensorboard} log to tensorboard, {learning_rate} learning rate, {weight_decay} weight decay, {seed} seed, {output_dir} output dir")
             iseed = seed
             model = FastLanguageModel.get_peft_model(
                 model,
@@ -300,7 +326,7 @@ with gr.Blocks(title="Unsloth fine-tuning") as demo:
             return "Model trained 100%",gr.update(visible=True, interactive=False), gr.update(visible=True, interactive=True), gr.update(interactive=True)
 
 
-        train_btn.click(train_model, inputs=[model_name, lora_r, lora_alpha, lora_dropout, per_device_train_batch_size, warmup_steps, max_steps, gradient_accumulation_steps, logging_steps, log_to_tensorboard, optim, learning_rate, weight_decay, lr_scheduler_type, seed, output_dir], outputs=[train_output, train_btn])
+        train_btn.click(train_model, inputs=[model_name, lora_r, lora_alpha, lora_dropout, per_device_train_batch_size, warmup_steps, max_steps, gradient_accumulation_steps, logging_steps, log_to_tensorboard, learning_rate, weight_decay, seed, output_dir], outputs=[train_output, train_btn])
 
     with gr.Tab("Save & Push Options"):