Hugging Face's logo

Spaces:
almamunkhan
/
MamunAI2
Running on Zero

App Files Community
MamunAI2
File size: 4,449 Bytes
9aab617
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
import gradio as gr
import spaces
import subprocess
import threading
import os

STATUS = "Idle"


# ── Training ──────────────────────────────────────────────────────────────────

@spaces.GPU
def run_training():
    global STATUS
    STATUS = "Training Running..."

    try:
        subprocess.run(
            ["python", "train.py"],
            check=True
        )
        STATUS = "Training Complete ✅"
    except Exception as e:
        STATUS = f"Training Failed ❌\n{e}"

    return STATUS


def start_training():
    thread = threading.Thread(target=run_training, daemon=False)
    thread.start()
    return "GPU requested. Training will start when GPU is allocated."


# ── Merge ─────────────────────────────────────────────────────────────────────


def run_merge(upload_repo: str):
    global STATUS
    STATUS = "Merging..."

    cmd = ["python", "train.py", "--merge"]
    if upload_repo.strip():
        cmd += ["--upload-merged-to", upload_repo.strip()]

    try:
        subprocess.run(cmd, check=True)
        STATUS = "Merge Complete ✅"
    except Exception as e:
        STATUS = f"Merge Failed ❌\n{e}"

    return STATUS


def start_merge(upload_repo: str):
    thread = threading.Thread(
        target=run_merge, args=(upload_repo,), daemon=False
    )
    thread.start()
    return "GPU requested. Merge will start when GPU is allocated."


# ── Status ────────────────────────────────────────────────────────────────────

def get_status():
    global STATUS
    return STATUS


# ── UI ────────────────────────────────────────────────────────────────────────

with gr.Blocks(title="MamunAI Trainer") as demo:

    gr.Markdown(
        """
        # 🤖 MamunAI Trainer
        **Owner:** Al Mamun Khan  |  **Space:** almamunkhan/MamunAI2

        Space is online.
        Press **Start Training** once — the server handles the rest.
        Training always resumes from the latest adapter in `almamunkhan/MamunAI`.
        Dataset is downloaded from GitHub automatically before every run.
        """
    )

    # ── Training section ──────────────────────────────────────────────────────
    gr.Markdown("---")
    gr.Markdown("## ▶ Training")

    start_btn = gr.Button("🚀 Start Training", variant="primary")

    status_box = gr.Textbox(label="Status", value="Idle")

    refresh_btn = gr.Button("🔄 Refresh Status")

    # ── Merge section ─────────────────────────────────────────────────────────
    gr.Markdown("---")
    gr.Markdown(
        "## 🔀 Merge Model\n"
        "_Merges the saved LoRA adapter into the base model. "
        "Only run this when you want the final merged weights — "
        "training never merges automatically._"
    )

    upload_repo_box = gr.Textbox(
        label="Upload merged model to HF repo (optional)",
        placeholder="e.g. almamunkhan/MamunAI-merged  — leave blank to skip",
    )

    merge_btn = gr.Button("🔀 Merge Model", variant="secondary")

    merge_status_box = gr.Textbox(label="Merge Status", value="Idle")

    merge_refresh_btn = gr.Button("🔄 Refresh Merge Status")

    # ── Event wiring ──────────────────────────────────────────────────────────
    start_btn.click(fn=start_training, outputs=status_box)
    refresh_btn.click(fn=get_status, outputs=status_box)

    merge_btn.click(fn=start_merge, inputs=upload_repo_box, outputs=merge_status_box)
    merge_refresh_btn.click(fn=get_status, outputs=merge_status_box)


if __name__ == "__main__":
    demo.queue()
    demo.launch(server_name="0.0.0.0")