show loss/epoch chart on finetune ui

llama_lora/ui/finetune/finetune_ui.py

@@ -28,7 +28,8 @@ from .previewing import (
 )
 from .training import (
     do_train,
-    render_training_status
+    render_training_status,
+    render_loss_plot
 )
 
 register_css_style('finetune', relative_read_file(__file__, "style.css"))
@@ -773,10 +774,15 @@ def finetune_ui():
         )
 
         train_status = gr.HTML(
-            "Training results will be shown here.",
+            "",
             label="Train Output",
             elem_id="finetune_training_status")
 
+        with gr.Column(visible=False, elem_id="finetune_loss_plot_container") as loss_plot_container:
+            loss_plot = gr.Plot(
+                visible=False, show_label=False,
+                elem_id="finetune_loss_plot")
+
         training_indicator = gr.HTML(
             "training_indicator", visible=False, elem_id="finetune_training_indicator")
 
@@ -787,7 +793,8 @@ def finetune_ui():
                 continue_from_model,
                 continue_from_checkpoint,
             ]),
-            outputs=[train_status, training_indicator]
+            outputs=[train_status, training_indicator,
+                     loss_plot_container, loss_plot]
         )
 
         # controlled by JS, shows the confirm_abort_button
@@ -803,6 +810,12 @@ def finetune_ui():
         outputs=[train_status, training_indicator],
         every=0.2
     )
+    loss_plot_updates = finetune_ui_blocks.load(
+        fn=render_loss_plot,
+        inputs=None,
+        outputs=[loss_plot_container, loss_plot],
+        every=10
+    )
     finetune_ui_blocks.load(_js=relative_read_file(__file__, "script.js"))
 
     # things_that_might_timeout.append(training_status_updates)

llama_lora/ui/finetune/style.css

@@ -255,7 +255,9 @@
     display: none;
 }
 
-#finetune_training_status > .wrap {
+#finetune_training_status > .wrap,
+#finetune_loss_plot_container > .wrap,
+#finetune_loss_plot > .wrap {
     border: 0;
     background: transparent;
     pointer-events: none;
@@ -264,6 +266,17 @@
     left: 0;
     right: 0;
 }
+#finetune_training_status > .wrap:not(.generating)::after {
+    content: "Refresh the page if this takes too long.";
+    position: absolute;
+    top: 0;
+    left: 0;
+    right: 0;
+    bottom: 0;
+    padding-top: 64px;
+    opacity: 0.5;
+    text-align: center;
+}
 #finetune_training_status > .wrap .meta-text-center {
     transform: none !important;
 }
@@ -383,5 +396,18 @@
     /* background: var(--error-background-fill) !important; */
     border: 1px solid var(--error-border-color) !important;
 }
+#finetune_loss_plot {
+    padding: var(--block-padding);
+}
+#finetune_loss_plot .altair {
+    overflow: auto !important;
+}
+#finetune_loss_plot .altair > * {
+    margin: auto !important;
+}
+#finetune_loss_plot .vega-embed summary {
+    border: 0;
+    box-shadow: none;
+}
 
 #finetune_training_indicator { display: none; }

llama_lora/ui/finetune/training.py

@@ -1,11 +1,14 @@
 import os
 import json
 import time
+import math
 import datetime
 import pytz
 import socket
 import threading
 import traceback
+import altair as alt
+import pandas as pd
 import gradio as gr
 
 from huggingface_hub import try_to_load_from_cache, snapshot_download
@@ -71,7 +74,7 @@ def do_train(
     progress=gr.Progress(track_tqdm=False),
 ):
     if Global.is_training or Global.is_train_starting:
-        return render_training_status()
+        return render_training_status() + render_loss_plot()
 
     reset_training_status()
     Global.is_train_starting = True
@@ -206,6 +209,9 @@ def do_train(
                     print(message)
 
                     total_steps = 300
+                    log_history = []
+                    initial_loss = 2
+                    loss_decay_rate = 0.8
                     for i in range(300):
                         if (Global.should_stop_training):
                             break
@@ -213,11 +219,14 @@ def do_train(
                         current_step = i + 1
                         total_epochs = 3
                         current_epoch = i / 100
-                        log_history = []
 
                         if (i > 20):
-                            loss = 3 + (i - 0) * (0.5 - 3) / (300 - 0)
-                            log_history = [{'loss': loss}]
+                            loss = initial_loss * math.exp(-loss_decay_rate * current_epoch)
+                            log_history.append({
+                                'loss': loss,
+                                'learning_rate': 0.0001,
+                                'epoch': current_epoch
+                            })
 
                         update_training_states(
                             total_steps=total_steps,
@@ -295,7 +304,7 @@ def do_train(
     finally:
         Global.is_train_starting = False
 
-    return render_training_status()
+    return render_training_status() + render_loss_plot()
 
 
 def render_training_status():
@@ -411,6 +420,51 @@ def render_training_status():
     return (gr.HTML.update(value=html_content), gr.HTML.update(visible=True))
 
 
+def render_loss_plot():
+    if len(Global.training_log_history) <= 2:
+        return (gr.Column.update(visible=False), gr.Plot.update(visible=False))
+
+    training_log_history = Global.training_log_history
+
+    loss_data = [
+        {
+            'type': 'train_loss' if 'loss' in item else 'eval_loss',
+            'loss': item.get('loss') or item.get('eval_loss'),
+            'epoch': item.get('epoch')
+        } for item in training_log_history
+        if ('loss' in item or 'eval_loss' in item)
+        and 'epoch' in item
+    ]
+
+    source = pd.DataFrame(loss_data)
+
+    highlight = alt.selection(
+        type='single',  # type: ignore
+        on='mouseover', fields=['type'], nearest=True
+    )
+
+    base = alt.Chart(source).encode(  # type: ignore
+        x='epoch:Q',
+        y='loss:Q',
+        color='type:N',
+        tooltip=['type:N', 'loss:Q', 'epoch:Q']
+    )
+
+    points = base.mark_circle().encode(
+        opacity=alt.value(0)
+    ).add_selection(
+        highlight
+    ).properties(
+        width=640
+    )
+
+    lines = base.mark_line().encode(
+        size=alt.condition(~highlight, alt.value(1), alt.value(3))
+    )
+
+    return (gr.Column.update(visible=True), gr.Plot.update(points + lines, visible=True))
+
+
 def format_time(seconds):
     hours, remainder = divmod(seconds, 3600)
     minutes, seconds = divmod(remainder, 60)

requirements.txt

@@ -1,4 +1,5 @@
 accelerate
+altair
 appdirs
 bitsandbytes
 black
@@ -7,10 +8,11 @@ datasets
 fire
 git+https://github.com/huggingface/peft.git
 git+https://github.com/huggingface/transformers.git
+gradio
 huggingface_hub
+loralib
 numba
 nvidia-ml-py3
-gradio
-loralib
-sentencepiece
+pandas
 random-word
+sentencepiece