Update app.py

app.py

@@ -1,96 +1,83 @@
 import os
-import time
-import zipfile
-import gradio as gr
+import shutil
 import subprocess
-import threading
-
-# Paths
-OUTPUT_DIR = "train_output"
-ZIP_FILE = "python_ai_trained_model.zip"
-LOG_FILE = "train_log.txt"
-
-# Zip function
-def zip_trained_model():
-    with zipfile.ZipFile(ZIP_FILE, 'w', zipfile.ZIP_DEFLATED) as zipf:
-        for root, _, files in os.walk(OUTPUT_DIR):
-            for file in files:
-                filepath = os.path.join(root, file)
-                arcname = os.path.relpath(filepath, OUTPUT_DIR)
-                zipf.write(filepath, arcname)
-
-# Tail logs function
-def tail_logs(n=20):
-    if not os.path.exists(LOG_FILE):
-        return ""
-    with open(LOG_FILE, 'r') as f:
-        return ''.join(f.readlines()[-n:])
-
-# Background training runner
-def run_training(status_box, time_box, download_file, model_size_box, log_box):
-    start_time = time.time()
-    status_box.value = "🚀 Training started..."
-    time_box.value = "Training in progress..."
-    log_box.value = ""
-
-    # Create log file
-    with open(LOG_FILE, "w") as log:
-        log.write("🚀 Launching train.py...\n")
-
-    # Start training
-    with open(LOG_FILE, "a") as log:
-        process = subprocess.Popen(["python", "train.py"], stdout=log, stderr=subprocess.STDOUT)
-        while process.poll() is None:
-            log_box.update(value=tail_logs())
-            time.sleep(5)
-
-    # Check exit status
-    if process.returncode != 0:
-        status_box.value = f"❌ Training failed with exit code {process.returncode}"
-        log_box.value = tail_logs()
-        return
-
-    # Training success
-    elapsed = round(time.time() - start_time, 2)
-    time_box.value = f"✅ Completed in {elapsed // 60:.0f} min {elapsed % 60:.0f} sec"
-    status_box.value = "🔄 Compressing model..."
-
-    # Zip it
-    zip_trained_model()
-    size_mb = round(os.path.getsize(ZIP_FILE) / (1024 * 1024), 2)
-    model_size_box.value = f"📦 Model Size: {size_mb} MB"
-
-    download_file.value = ZIP_FILE
-    download_file.visible = True
-    status_box.value = "✅ Training complete. Download below."
-
-# Button trigger
-def start_training(status_box, time_box, download_file, model_size_box, log_box):
-    thread = threading.Thread(target=run_training, args=(status_box, time_box, download_file, model_size_box, log_box))
-    thread.start()
-    return "Training process started."
-
-# Gradio UI
-with gr.Blocks() as demo:
-    gr.Markdown("## 🧠 Python AI Trainer (StarCoder 7B)")
-    gr.Markdown("Train your Python AI with 1 click. Watch logs. Download model when done.")
-
-    with gr.Row():
-        train_btn = gr.Button("🚀 Start Training")
-        status_box = gr.Textbox(label="Status", value="Ready", interactive=False)
-
-    with gr.Row():
-        time_box = gr.Textbox(label="Training Time", interactive=False)
-        model_size_box = gr.Textbox(label="Final Model Size", interactive=False)
-
-    log_box = gr.Textbox(label="Live Training Logs", lines=20, interactive=False, value="")
-    download_file = gr.File(label="📥 Download Trained Model (.zip)", visible=False)
-
-    train_btn.click(
-        fn=start_training,
-        inputs=[status_box, time_box, download_file, model_size_box, log_box],
-        outputs=[status_box]
-    )
-
-    
-demo.launch()
+import gradio as gr
+from transformers import pipeline
+
+uploaded_file_path = "dataset.jsonl"
+log_path = "train.log"
+model_dir = "trained_model"
+zip_file = "trained_model.zip"
+
+# Try loading the generator model once at the top (for better performance)
+try:
+    generator = pipeline("text-generation", model=model_dir, tokenizer="distilgpt2")
+except:
+    generator = None
+
+def upload_file(file):
+    if file is None:
+        return "No file uploaded."
+    shutil.copy(file.name, uploaded_file_path)
+    return "✅ File uploaded successfully."
+
+def start_training():
+    with open(log_path, "w") as log_file:
+        process = subprocess.Popen(
+            ["python", "train.py", "--dataset", uploaded_file_path],
+            stdout=log_file,
+            stderr=subprocess.STDOUT
+        )
+        process.wait()
+
+    if os.path.exists(model_dir):
+        shutil.make_archive("trained_model", "zip", model_dir)
+        return "✅ Training complete!", zip_file
+    else:
+        return "❌ Training failed.", None
+
+def read_logs():
+    if os.path.exists(log_path):
+        with open(log_path, "r") as f:
+            return f.read()
+    return "Waiting for logs..."
+
+def generate_response(prompt):
+    try:
+        if generator is None:
+            return "❌ Model not loaded. Please train or upload a valid model."
+        result = generator(
+            prompt,
+            max_length=256,
+            do_sample=True,
+            temperature=0.7,
+            truncation=True  # ✅ Fix warning and enforce consistent length
+        )[0]["generated_text"]
+        return result
+    except Exception as e:
+        return f"❌ Error: {e}"
+
+# === UI ===
+with gr.Blocks() as app:
+    with gr.Tab("🧠 Train AI"):
+        gr.Markdown("## 📥 Upload your dataset and 🎯 train a Godot AI")
+        file_input = gr.File(label="Upload JSONL Dataset")
+        upload_btn = gr.Button("Upload")
+        status_box = gr.Textbox(label="Upload Status")
+
+        start_btn = gr.Button("🚀 Start Training")
+        log_output = gr.Textbox(label="📜 Training Logs", lines=15)
+        download_btn = gr.File(label="📥 Download Trained Model", visible=False)
+
+        upload_btn.click(fn=upload_file, inputs=file_input, outputs=status_box)
+        start_btn.click(fn=start_training, outputs=[status_box, download_btn])
+        start_btn.click(fn=read_logs, outputs=log_output)
+
+    with gr.Tab("🚀 Test AI"):
+        gr.Markdown("## 💡 Try your trained Godot AI below")
+        prompt_input = gr.Textbox(label="Enter Prompt")
+        test_btn = gr.Button("🔍 Test AI")
+        response_output = gr.Textbox(label="AI Response", lines=10)
+        test_btn.click(fn=generate_response, inputs=prompt_input, outputs=response_output)
+
+app.launch()