Update app_gpu.py

app_gpu.py

@@ -291,7 +291,7 @@ def run_ui():
                       outputs=[logs],
                       queue=True)
 
-        # ---------------- Inference Tab ----------------
+        # ---------------- Inference (CPU) Tab ----------------
         with gr.Tab("Inference (CPU)"):
             inf_base_model = gr.Textbox(label="Base model", value="google/gemma-3-4b-it")
             inf_lora_repo = gr.Textbox(label="LoRA HF repo", value="rahul7star/gemma-3-270m-ccebc0")
@@ -305,63 +305,53 @@ def run_ui():
 
         # ---------------- Code Explain Tab ----------------
         with gr.Tab("Code Explain"):
-
-            def simulate_logs(base_model, r_, a_, ep_):
-                simulated = [
-                    f"[INFO] Loading base model: {base_model}",
-                    f"[INFO] LoRA configuration: r={r_}, alpha={a_}",
-                    f"[INFO] Epoch {i+1}/{ep_} started..." for i in range(int(ep_))
-                ]
-                for ep_idx in range(int(ep_)):
-                    for step in range(1, 6):
-                        simulated.append(f"[DEBUG] Step {step}, Loss: {0.01 * (6-step):.6f}")
-                    simulated.append(f"[INFO] Epoch {ep_idx+1} completed.")
-                simulated.append("[INFO] LoRA training finished. Ready to upload to HF Hub.")
-                return "\n".join(simulated)
-
-            model_explain = gr.Textbox(label="Base Model", value="google/gemma-3-4b-it")
-            lora_rank = gr.Number(label="LoRA rank (r)", value=8)
-            lora_alpha = gr.Number(label="LoRA alpha", value=16)
-            epochs = gr.Number(label="Epochs", value=1)
-            logs_out = gr.Textbox(label="Simulated Logs & Explanation", lines=30)
-
-            def explain_code(model_name, r_, a_, ep_):
-                logs = simulate_logs(model_name, r_, a_, ep_)
-                explanation = f"""
-### Universal LoRA Trainer & Inference - Detailed Explanation
-
-1. **Imports**: Handles data, tensor ops, LoRA PEFT, HF Hub integration, and optional Transformers.
-2. **Dataset**: `MediaTextDataset` loads short & long prompts from CSV/Parquet/HF.
-3. **Model Loader**: Loads base Gemma model; detects Linear layers (Q/K/V) to apply LoRA.
-4. **LoRA Internals**:
-   - LoRA injects low-rank matrices `A` and `B` into Q/K/V projections.
-   - `Effective weight: W_eff = W + alpha * B @ A`
-   - Only LoRA parameters are trained; main model frozen.
-5. **Training Loop**:
-   - Forward pass → Cross-entropy loss.
-   - Backprop updates LoRA weights only.
-   - Accelerator handles device placement & mixed precision.
-6. **CPU Inference**:
-   - Loads base + LoRA on CPU.
-   - Merges LoRA optionally to avoid runtime PEFT issues.
-   - Generates expanded prompt from short prompt.
-7. **Gradio UI Tabs**:
-   - Train LoRA: Configure training, see live logs.
-   - Inference: Expand short prompts using LoRA.
-   - Code Explain: This simulation showing internal workflow & parameter effects.
-
-**Simulated Training Logs:**\n{logs}
-"""
-                return explanation
-
-            explain_btn = gr.Button("📝 Show Code Explanation & Logs")
-            explain_btn.click(fn=explain_code,
-                              inputs=[model_explain, lora_rank, lora_alpha, epochs],
-                              outputs=[logs_out])
+            explain_md = gr.Markdown("""
+### Universal LoRA Trainer & Inference - Code Explanation
+
+#### 1. Imports
+- `spaces, os, torch, gradio, pandas, numpy`: utilities, tensor ops, UI, and data handling.
+- `peft (LoraConfig, get_peft_model)`: LoRA adapter integration.
+- `accelerate (Accelerator)`: device placement, mixed precision, distributed training.
+- `huggingface_hub`: upload LoRA weights to Hugging Face.
+- `transformers (optional)`: only for Gemma LLM.
+
+#### 2. Dataset
+- `MediaTextDataset`: Loads CSV/Parquet or HF dataset, extracts `short_prompt` and `long_prompt`.
+- Handles batched access and missing columns.
+
+#### 3. Model Loading
+- `load_pipeline_auto`: Loads Gemma tokenizer + model in float16/32.
+- `find_target_modules`: Detects Linear layers (Q/K/V) for LoRA injection.
+
+#### 4. LoRA Training
+- LoRA formula: `W_eff = W + alpha * B @ A`
+- `r` = low-rank dimension, `alpha` = scaling factor.
+- Only trains LoRA matrices, main model frozen.
+- Efficient memory & compute, streams logs.
+- Supports uploading trained LoRA to Hugging Face Hub.
+
+#### 5. CPU Inference
+- Loads base Gemma model on CPU (float32).
+- Loads LoRA with `PeftModel.from_pretrained`.
+- Optionally merges LoRA into base model.
+- Generates long prompt using `generate()` with top-p/top-k sampling.
+
+#### 6. LoRA Internals
+- LoRA injects trainable matrices `A` & `B` into selected Linear layers.
+- Q/K/V matrices in attention updated as `Q_new = Q + alpha*B@A`.
+- Efficient: `r << hidden_size`, only small matrices trained.
+
+#### 7. Gradio UI
+- Train Tab: configure model, dataset, LoRA, HF repo.
+- Inference Tab: short prompt → expanded long prompt.
+- Code Explain Tab: shows detailed explanation and simulated logs.
+""")
+            explain_md.render()
 
     return demo
 
 
+