Update app.py

app.py

@@ -1,168 +1,186 @@
+# ==============================================================
+#                       EMOTION CLASSIFIER
+# ==============================================================
+
 import os
-import io
 import math
 import torch
 import pandas as pd
 import numpy as np
-import matplotlib.pyplot as plt
 import gradio as gr
+import matplotlib.pyplot as plt
+
 from pathlib import Path
 from torch import nn
 from torch.utils.data import Dataset, DataLoader, TensorDataset
 from sklearn.model_selection import train_test_split
-from transformers import AutoTokenizer, AutoModel, get_linear_schedule_with_warmup
-
-# ---------------------------
-# Konfigurasi & Label
-# ---------------------------
+from transformers import (
+    AutoTokenizer, 
+    AutoModel, 
+    AutoConfig,
+    get_linear_schedule_with_warmup
+)
+
+# =========================================================
+# CONFIG
+# =========================================================
 LABELS = ['anger','anticipation','disgust','fear','joy','sadness','surprise','trust']
 LABEL2ID = {l:i for i,l in enumerate(LABELS)}
 ID2LABEL = {i:l for i,l in enumerate(LABELS)}
+
 SAVED_ROOT = Path("saved_models")
 SAVED_ROOT.mkdir(exist_ok=True)
 
-# ---------------------------
-# Utility I/O small helpers
-# ---------------------------
+# ==============================================================
+#                     Simpan dan Muat Data
+# ==============================================================
 def read_uploaded_file(uploaded):
-    # uploaded can be a gradio file object or a path string
     if uploaded is None:
         raise ValueError("No file provided")
+
     if isinstance(uploaded, str):
         return uploaded
-    # gradio returns a tempfile-like object with 'name' attribute
+
     if hasattr(uploaded, "name"):
         return uploaded.name
-    # fallback: bytesIO-like
+
     if hasattr(uploaded, "read"):
-        # write to temp file
         tmp = Path("/tmp") / f"uploaded_{np.random.randint(1e9)}.csv"
         with open(tmp, "wb") as f:
             f.write(uploaded.read())
         return str(tmp)
+
     raise ValueError("Unsupported uploaded file type")
 
-def save_last_model_name(model_name: str):
-    (SAVED_ROOT / "last_model.txt").write_text(model_name)
 
-def load_last_model_name() -> str:
+def save_last_model_name(name):
+    (SAVED_ROOT / "last_model.txt").write_text(name)
+
+
+def load_last_model_name():
     p = SAVED_ROOT / "last_model.txt"
     if p.exists():
         return p.read_text().strip()
     return None
 
-def model_folder(model_name: str) -> Path:
+
+def model_folder(model_name):
     return SAVED_ROOT / model_name.replace("/", "_")
 
-# ---------------------------
-# Data loading & cleaning
-# ---------------------------
-def load_dataset(path_or_file, sep=","):
-    path = read_uploaded_file(path_or_file)
-    df = pd.read_csv(path, sep=sep)
-    return df
 
+# ==============================================================
+#                        Pembersihan Data
+# ==============================================================
 def clean_labels(df):
-    # ensure all LABELS exist as columns (0/1)
     for l in LABELS:
         if l not in df.columns:
             df[l] = 0
     return df
 
-def clean_text(df, text_col="text"):
-    if text_col not in df.columns:
-        raise KeyError(f"CSV must contain column named '{text_col}' (found columns: {df.columns.tolist()})")
-    df[text_col] = df[text_col].astype(str).str.replace("\n", " ").str.strip()
+
+def clean_text(df, col="text"):
+    if col not in df.columns:
+        raise KeyError(f"CSV must contain a column '{col}'")
+    df[col] = df[col].astype(str).str.replace("\n", " ").str.strip()
     return df
 
-# ---------------------------
-# Model class (BERT + head)
-# ---------------------------
-class EmotionClassifier(nn.Module):
-    def __init__(self, model_name="bert-base-multilingual-cased", num_labels=8):
+
+# =========================================================
+#                        Model AI
+# =========================================================
+class EmotionModel(nn.Module):
+    """Consistent backbone + dropout + classifier."""
+    def __init__(self, base_model_name, num_labels=8):
         super().__init__()
-        self.bert = AutoModel.from_pretrained(model_name)
+        self.config = AutoConfig.from_pretrained(base_model_name)
+        self.base = AutoModel.from_pretrained(base_model_name)
         self.drop = nn.Dropout(0.3)
-        self.classifier = nn.Linear(self.bert.config.hidden_size, num_labels)
-
-    def forward(self, input_ids, attention_mask=None, token_type_ids=None):
-        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)
-        cls = outputs.last_hidden_state[:,0,:]
-        x = self.drop(cls)
-        logits = self.classifier(x)
-        return logits
-
-# ---------------------------
-# Tokenisasi dan dataset (optimized batch)
-# ---------------------------
-def tokenize_dataset_batch(texts, tokenizer, max_len=128):
-    enc = tokenizer.batch_encode_plus(
+        self.clf = nn.Linear(self.config.hidden_size, num_labels)
+
+    def forward(self, ids, mask):
+        out = self.base(
+            input_ids=ids,
+            attention_mask=mask
+        )
+
+        # Prefer pooler_output if exists
+        if hasattr(out, "pooler_output") and out.pooler_output is not None:
+            x = out.pooler_output
+        else:
+            x = out.last_hidden_state[:, 0, :]
+
+        x = self.drop(x)
+        return self.clf(x)
+
+
+# ==============================================================
+#                       Tokenisasi Dataset
+# ==============================================================
+def tokenize_batch(texts, tokenizer, max_len=128):
+    return tokenizer(
         texts,
         padding="max_length",
         truncation=True,
         max_length=max_len,
         return_tensors="pt"
     )
-    return enc  # dict: input_ids, attention_mask, (token_type_ids)
+
 
 def build_tensor_dataset(df, tokenizer, max_len=128):
-    texts = df["text"].tolist()
-    enc = tokenize_dataset_batch(texts, tokenizer, max_len=max_len)
+    enc = tokenize_batch(df["text"].tolist(), tokenizer, max_len)
     labels = torch.tensor(df[LABELS].values, dtype=torch.float)
-    dataset = TensorDataset(enc["input_ids"], enc["attention_mask"], labels)
-    return dataset
+    return TensorDataset(
+        enc["input_ids"],
+        enc["attention_mask"],
+        labels
+    )
+
 
-# ---------------------------
-# Pos-weight compute
-# ---------------------------
+# ==============================================================
+#                           Bobot
+# ==============================================================
 def compute_pos_weight(df):
-    counts = df[LABELS].sum(axis=0).astype(int).to_list()
+    counts = df[LABELS].sum(axis=0)
     N = len(df)
     pw = []
     for c in counts:
-        if c == 0:
-            pw.append(1.0)
-        else:
-            pw.append((N - c) / c)
+        pw.append((N - c) / c if c > 0 else 1.0)
     return torch.tensor(pw, dtype=torch.float)
 
-# ---------------------------
-# Save / Load trained model files
-# ---------------------------
-def save_trained(model, tokenizer, model_name:str):
-    folder = model_folder(model_name)
-    folder.mkdir(parents=True, exist_ok=True)
-    # save model weights
-    torch.save(model.state_dict(), folder / "best_model.pt")
-    # save tokenizer config
-    tokenizer.save_pretrained(str(folder))
-    # save a text marker
-    save_last_model_name(model_name)
-    return str(folder)
-
-def load_trained(model_name: str = None, device=None):
-    if model_name is None:
-        model_name = load_last_model_name()
-        if model_name is None:
-            raise ValueError("No trained model found. Train a model first.")
-    folder = model_folder(model_name)
-    if not folder.exists():
-        raise FileNotFoundError(f"Saved model folder not found: {folder}")
-    # load tokenizer and instantiate model then load state dict
-    tokenizer = AutoTokenizer.from_pretrained(str(folder))
-    # we need the original base model identifier to instantiate architecture.
-    # Assume original model_name saved in folder name; instantiate using folder's config via AutoModel? We used AutoModel, but for simplicity re-use model_name prefix by reading tokenizer._name_or_path if possible
-    base_name = tokenizer.name_or_path if hasattr(tokenizer, "name_or_path") else model_name
-    model = EmotionClassifier(base_name)
-    state = torch.load(folder / "best_model.pt", map_location=device or "cpu")
-    model.load_state_dict(state)
-    if device:
-        model.to(device)
-    return model, tokenizer, model_name
-
-# ---------------------------
-# Training loop (uses trainable params only)
-# ---------------------------
+
+# ==============================================================
+#                       Simpan dan Muat Model
+# ==============================================================
+def save_model(model, tokenizer, folder):
+    os.makedirs(folder, exist_ok=True)
+
+    # Save backbone HF style
+    model.base.save_pretrained(folder)
+    tokenizer.save_pretrained(folder)
+
+    # Save classifier head
+    torch.save(model.clf.state_dict(), str(Path(folder) / "classifier.pt"))
+
+    # Save last-used name
+    save_last_model_name(str(folder))
+
+
+def load_model(folder):
+    folder = str(folder)
+    config = AutoConfig.from_pretrained(folder)
+    tokenizer = AutoTokenizer.from_pretrained(folder)
+
+    model = EmotionModel(folder)
+    state = torch.load(f"{folder}/classifier.pt", map_location="cpu")
+    model.clf.load_state_dict(state)
+    model.eval()
+
+    return model, tokenizer, config
+
+
+# ==============================================================
+#                            Pelatihan
+# ==============================================================
 def train_model(
     df,
     model_name="bert-base-multilingual-cased",
@@ -176,33 +194,34 @@ def train_model(
     freeze_layers=6,
     device=None
 ):
-    device = device or (torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu"))
+    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+
     tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-    # prepare dataset
     df = df.reset_index(drop=True)
-    enc_dataset = build_tensor_dataset(df, tokenizer, max_len=max_len)
-    # split indices
-    n = len(enc_dataset)
-    idx = list(range(n))
+    dataset = build_tensor_dataset(df, tokenizer, max_len)
+
+    idx = list(range(len(dataset)))
     train_idx, val_idx = train_test_split(idx, test_size=0.15, random_state=42)
-    def subset(ds, indices):
-        input_ids = torch.stack([ds[i][0] for i in indices])
-        attn = torch.stack([ds[i][1] for i in indices])
-        labels = torch.stack([ds[i][2] for i in indices])
-        return TensorDataset(input_ids, attn, labels)
 
-    train_ds = subset(enc_dataset, train_idx)
-    val_ds = subset(enc_dataset, val_idx)
+    def subset(ds, idxs):
+        return TensorDataset(
+            torch.stack([ds[i][0] for i in idxs]),
+            torch.stack([ds[i][1] for i in idxs]),
+            torch.stack([ds[i][2] for i in idxs]),
+        )
+
+    train_ds = subset(dataset, train_idx)
+    val_ds = subset(dataset, val_idx)
 
     train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True)
     val_loader = DataLoader(val_ds, batch_size=batch_size)
 
-    model = EmotionClassifier(model_name)
+    model = EmotionModel(model_name)
     model.to(device)
 
-    # freeze layers if requested (works for BERT-like named params)
-    for name, param in model.bert.named_parameters():
+    # Freeze lower layers
+    for name, param in model.base.named_parameters():
         if name.startswith("embeddings."):
             param.requires_grad = False
         elif name.startswith("encoder.layer"):
@@ -216,127 +235,154 @@ def train_model(
     pos_weight = compute_pos_weight(df).to(device)
     loss_fn = nn.BCEWithLogitsLoss(pos_weight=pos_weight)
 
-    optimizer = torch.optim.AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=lr, weight_decay=weight_decay)
+    optimizer = torch.optim.AdamW(
+        filter(lambda p: p.requires_grad, model.parameters()),
+        lr=lr,
+        weight_decay=weight_decay
+    )
+
     total_steps = len(train_loader) * epochs
-    warmup_steps = int(warmup_ratio * total_steps) if total_steps>0 else 0
-    scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_steps)
+    warmup_steps = int(warmup_ratio * total_steps)
+
+    scheduler = get_linear_schedule_with_warmup(
+        optimizer,
+        num_warmup_steps=warmup_steps,
+        num_training_steps=total_steps
+    )
 
     best_val = float("inf")
-    epochs_no_improve = 0
-    history = {"train_loss":[], "val_loss":[]}
+    no_improve = 0
+
+    history = {"train_loss": [], "val_loss": []}
+
+    save_path = str(model_folder(model_name))
 
-    for epoch in range(1, epochs+1):
+    for ep in range(1, epochs+1):
         model.train()
-        running_loss = 0.0
-        for batch in train_loader:
-            optimizer.zero_grad()
-            input_ids = batch[0].to(device)
-            attn = batch[1].to(device)
-            labels = batch[2].to(device)
+        t_loss = 0
+
+        for input_ids, attn, labels in train_loader:
+            input_ids = input_ids.to(device)
+            attn = attn.to(device)
+            labels = labels.to(device)
 
-            logits = model(input_ids=input_ids, attention_mask=attn)
+            optimizer.zero_grad()
+            logits = model(input_ids, attn)
             loss = loss_fn(logits, labels)
             loss.backward()
             optimizer.step()
-            if scheduler is not None:
-                scheduler.step()
+            scheduler.step()
 
-            running_loss += loss.item() * input_ids.size(0)
+            t_loss += loss.item() * input_ids.size(0)
 
-        avg_train = running_loss / len(train_loader.dataset)
-        history["train_loss"].append(avg_train)
+        train_loss = t_loss / len(train_loader.dataset)
+        history["train_loss"].append(train_loss)
 
-        # validation
+        # Validation
         model.eval()
-        vloss = 0.0
+        v_loss = 0
         with torch.no_grad():
-            for batch in val_loader:
-                input_ids = batch[0].to(device)
-                attn = batch[1].to(device)
-                labels = batch[2].to(device)
-                logits = model(input_ids=input_ids, attention_mask=attn)
+            for input_ids, attn, labels in val_loader:
+                input_ids = input_ids.to(device)
+                attn = attn.to(device)
+                labels = labels.to(device)
+                logits = model(input_ids, attn)
                 loss = loss_fn(logits, labels)
-                vloss += loss.item() * input_ids.size(0)
-        avg_val = vloss / len(val_loader.dataset)
-        history["val_loss"].append(avg_val)
-
-        print(f"Epoch {epoch}/{epochs} | Train loss {avg_train:.4f} | Val loss {avg_val:.4f}")
-
-        if avg_val < best_val:
-            best_val = avg_val
-            epochs_no_improve = 0
-            # save model+tokenizer to folder
-            save_trained(model, tokenizer, model_name)
-            print(f"Saved best model for {model_name}")
+                v_loss += loss.item() * input_ids.size(0)
+
+        val_loss = v_loss / len(val_loader.dataset)
+        history["val_loss"].append(val_loss)
+
+        print(f"Epoch {ep} | Train={train_loss:.4f} | Val={val_loss:.4f}")
+
+        if val_loss < best_val:
+            best_val = val_loss
+            no_improve = 0
+            save_model(model, tokenizer, save_path)
+            print(f"Saved best model to {save_path}")
         else:
-            epochs_no_improve += 1
-            if epochs_no_improve >= patience:
-                print("Early stopping triggered")
+            no_improve += 1
+            if no_improve >= patience:
+                print("Early stopping.")
                 break
 
     return model, tokenizer, history
 
-# ---------------------------
-# Inference helpers (batch optimized)
-# ---------------------------
-def predict_batch_from_texts(texts, model, tokenizer, max_len=128, batch_size=32, device=None):
-    device = device or (torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu"))
-    model.to(device)
-    model.eval()
-    results = []
-    # batch tokenize
+
+# ==============================================================
+#                              Uji
+# ==============================================================
+def predict_single(text, folder=None):
+    folder = folder or load_last_model_name()
+    model, tokenizer, cfg = load_model(folder)
+
+    encoded = tokenizer(
+        text,
+        padding="max_length",
+        truncation=True,
+        max_length=128,
+        return_tensors="pt"
+    )
+
+    with torch.no_grad():
+        out = model(encoded["input_ids"], encoded["attention_mask"])
+        probs = torch.sigmoid(out).numpy()[0]
+
+    return {LABELS[i]: float(probs[i]) for i in range(len(LABELS))}
+
+
+def predict_batch(texts, folder=None, batch_size=32):
+    folder = folder or load_last_model_name()
+    model, tokenizer, cfg = load_model(folder)
+
+    preds = []
     for i in range(0, len(texts), batch_size):
-        batch_texts = texts[i:i+batch_size]
-        enc = tokenizer.batch_encode_plus(batch_texts, padding="max_length", truncation=True, max_length=max_len, return_tensors="pt")
-        input_ids = enc["input_ids"].to(device)
-        attn = enc["attention_mask"].to(device)
+        batch = texts[i:i+batch_size]
+        enc = tokenizer(
+            batch,
+            padding="max_length",
+            truncation=True,
+            max_length=128,
+            return_tensors="pt"
+        )
+
         with torch.no_grad():
-            logits = model(input_ids=input_ids, attention_mask=attn)
-            probs = torch.sigmoid(logits).cpu().numpy()
+            out = model(enc["input_ids"], enc["attention_mask"])
+            probs = torch.sigmoid(out).numpy()
+
         for p in probs:
-            results.append({LABELS[j]: float(p[j]) for j in range(len(LABELS))})
-    return results
-
-def predict_single_using_saved(text, max_len=128, batch_size=32):
-    last = load_last_model_name()
-    if last is None:
-        raise ValueError("No trained model found. Train first.")
-    model, tokenizer, _ = load_trained(last)
-    res = predict_batch_from_texts([text], model, tokenizer, max_len=max_len, batch_size=batch_size)
-    return res[0]
-
-# ---------------------------
-# Summary utility
-# ---------------------------
-def summary_top3_from_preds(preds):
-    # preds: list of dict {label:prob}
-    avg = {l:0.0 for l in LABELS}
-    n = max(1, len(preds))
+            preds.append({LABELS[j]: float(p[j]) for j in range(len(LABELS))})
+
+    return preds
+
+
+def summarize_preds(preds):
+    avg = {l: 0.0 for l in LABELS}
+    n = len(preds)
+
     for p in preds:
         for l,v in p.items():
-            avg[l] += float(v)
+            avg[l] += v
     for l in avg:
         avg[l] /= n
-    sorted_avg = sorted(avg.items(), key=lambda x: x[1], reverse=True)
-    top3 = [{"label": sorted_avg[i][0], "score": float(sorted_avg[i][1])} for i in range(min(3, len(sorted_avg)))]
-    return {"n": n, "avg_distribution": avg, "top3": top3}
-
-# ---------------------------
-# Wrappers for GUI
-# ---------------------------
-def wrapper_training(
-    file_obj, sep=",",
-    model_name="bert-base-multilingual-cased",
-    epochs=3, batch_size=8, lr=2e-5, max_len=128,
-    weight_decay=0.01, warmup_ratio=0.1, patience=2, freeze_layers=6
-):
-    # file_obj can be gr.File or path string
-    csv_path = read_uploaded_file(file_obj)
-    df = pd.read_csv(csv_path, sep=sep)
+
+    top3 = sorted(avg.items(), key=lambda x: x[1], reverse=True)[:3]
+    top3 = [{"label":l, "score":float(s)} for l,s in top3]
+
+    return {"n":n, "avg_distribution":avg, "top3":top3}
+
+
+# ==============================================================
+#                             GRADIO GUI
+# ==============================================================
+def wrapper_train(file_obj, sep, model_name, epochs, batch_size, lr,
+                  max_len, weight_decay, warmup_ratio, patience, freeze_layers):
+    csv = read_uploaded_file(file_obj)
+    df = pd.read_csv(csv, sep=sep)
     df = clean_labels(df)
     df = clean_text(df)
 
-    model, tokenizer, history = train_model(
+    _, _, history = train_model(
         df=df,
         model_name=model_name,
         epochs=int(epochs),
@@ -348,86 +394,80 @@ def wrapper_training(
         patience=int(patience),
         freeze_layers=int(freeze_layers)
     )
-    # return a short report and history summary
+
     return {
-        "message": f"Training finished. Best model saved under saved_models/{model_name}",
-        "history": {"train_loss": history["train_loss"], "val_loss": history["val_loss"]},
+        "message": "Training finished.",
+        "history": history,
         "model_name": model_name
     }
 
-def wrapper_predict_single(text, max_len=128):
-    res = predict_single_using_saved(text, max_len=max_len)
-    return res
 
-def wrapper_predict_dataset(file_obj, sep=",", max_len=128, batch_size=32):
-    csv_path = read_uploaded_file(file_obj)
-    df = pd.read_csv(csv_path, sep=sep)
+def wrapper_single(text):
+    return predict_single(text)
+
+
+def wrapper_dataset(file_obj, sep, max_len, batch_size):
+    csv = read_uploaded_file(file_obj)
+    df = pd.read_csv(csv, sep=sep)
     df = clean_labels(df)
     df = clean_text(df)
-    texts = df["text"].tolist()
-    last = load_last_model_name()
-    if last is None:
-        return {"error":"No trained model found. Train first."}
-    model, tokenizer, _ = load_trained(last)
-    preds = predict_batch_from_texts(texts, model, tokenizer, max_len=max_len, batch_size=int(batch_size))
-    summary = summary_top3_from_preds(preds)
-    return {"n": summary["n"], "top3": summary["top3"], "avg_distribution": summary["avg_distribution"]}
-
-# ---------------------------
-# Plot helper (optional in notebook)
-# ---------------------------
-def plot_emotion_pie_from_avg(avg_dict):
-    labels = list(avg_dict.keys())
-    values = list(avg_dict.values())
-    plt.figure(figsize=(6,6))
-    plt.pie(values, labels=labels, autopct="%1.1f%%")
-    plt.title("Emotion Distribution (average)")
-    plt.show()
-
-# ---------------------------
-# Gradio GUI
-# ---------------------------
+
+    preds = predict_batch(df["text"].tolist(), batch_size=int(batch_size))
+    return summarize_preds(preds)
+
+
+# ==============================================================
+#                         Menjalankan GRADIO
+# ==============================================================
 with gr.Blocks() as app:
-    gr.Markdown("## Emotion Classifier — Dava (Revised)")
+    gr.Markdown("## Emotion Classifier — Dava (Final Version)")
 
     with gr.Tab("Training"):
-        file_in = gr.File(label="Upload training CSV")
+        file_in = gr.File(label="Upload Training CSV")
         sep_in = gr.Textbox(label="Delimiter", value=",")
-        model_name_in = gr.Dropdown(label="Model backbone", choices=[
-            "bert-base-multilingual-cased", "indobert-base-uncased", "bert-base-uncased"
-        ], value="bert-base-multilingual-cased")
+        model_name_in = gr.Dropdown(
+            label="Backbone Model",
+            choices=["bert-base-multilingual-cased", "indobert-base-p1"],
+            value="bert-base-multilingual-cased"
+        )
         epochs_in = gr.Number(label="Epochs", value=3)
-        batch_in = gr.Number(label="Batch size", value=8)
-        lr_in = gr.Number(label="Learning rate", value=2e-5)
-        maxlen_in = gr.Number(label="Max length", value=128)
-        weightdecay_in = gr.Number(label="Weight decay", value=0.01)
-        warmup_in = gr.Number(label="Warmup ratio", value=0.1)
-        patience_in = gr.Number(label="Early stop patience", value=2)
-        freeze_in = gr.Number(label="Freeze layers (first n)", value=6)
-        train_btn = gr.Button("Start Training")
-        train_out = gr.JSON(label="Training result (history + message)")
-
-        train_btn.click(
-            fn=wrapper_training,
-            inputs=[file_in, sep_in, model_name_in, epochs_in, batch_in, lr_in, maxlen_in, weightdecay_in, warmup_in, patience_in, freeze_in],
-            outputs=train_out
+        bs_in = gr.Number(label="Batch Size", value=8)
+        lr_in = gr.Number(label="Learning Rate", value=2e-5)
+        maxlen_in = gr.Number(label="Max Length", value=128)
+        wd_in = gr.Number(label="Weight Decay", value=0.01)
+        warmup_in = gr.Number(label="Warmup Ratio", value=0.1)
+        patience_in = gr.Number(label="Patience", value=2)
+        freeze_in = gr.Number(label="Freeze Layers", value=6)
+
+        btn_train = gr.Button("Start Training")
+        out_train = gr.JSON(label="Train Result")
+
+        btn_train.click(
+            wrapper_train,
+            inputs=[file_in, sep_in, model_name_in, epochs_in, bs_in,
+                    lr_in, maxlen_in, wd_in, warmup_in, patience_in, freeze_in],
+            outputs=out_train
         )
 
-    with gr.Tab("Single Inference"):
-        text_in = gr.Textbox(label="Text to analyze")
-        single_btn = gr.Button("Predict")
-        single_out = gr.JSON(label="Emotion probabilities")
-
-        single_btn.click(fn=wrapper_predict_single, inputs=[text_in], outputs=single_out)
+    with gr.Tab("Single Prediction"):
+        text_in = gr.Textbox(label="Text")
+        btn_single = gr.Button("Predict")
+        out_single = gr.JSON(label="Emotion Scores")
+        btn_single.click(wrapper_single, inputs=[text_in], outputs=out_single)
 
-    with gr.Tab("Dataset Inference"):
-        file_test = gr.File(label="Upload CSV for inference")
+    with gr.Tab("Dataset Prediction"):
+        file_test = gr.File(label="Upload CSV")
         sep_test = gr.Textbox(label="Delimiter", value=",")
-        maxlen_test = gr.Number(label="Max length", value=128)
-        batchsize_test = gr.Number(label="Batch size (inference)", value=32)
-        test_btn = gr.Button("Run Inference")
-        test_out = gr.JSON(label="Summary result")
+        maxlen_test = gr.Number(label="Max Length", value=128)
+        bs_test = gr.Number(label="Batch Size", value=32)
 
-        test_btn.click(fn=wrapper_predict_dataset, inputs=[file_test, sep_test, maxlen_test, batchsize_test], outputs=test_out)
+        btn_test = gr.Button("Run Prediction")
+        out_test = gr.JSON(label="Summary Result")
+
+        btn_test.click(
+            wrapper_dataset,
+            inputs=[file_test, sep_test, maxlen_test, bs_test],
+            outputs=out_test
+        )
 
-app.launch()
+app.launch()