app.py

import gradio as gr
import os
import subprocess
import sys
import threading
import time
from pathlib import Path

# Import your training code
from fine_tune_cove import train_cove_model, test_cove_inference

def install_requirements():
    """Install required packages"""
    try:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-r", "requirements.txt"])
        return "✅ Requirements installed successfully!"
    except Exception as e:
        return f"❌ Error installing requirements: {str(e)}"

def start_training(hf_token, model_name="codellama/CodeLlama-7b-Instruct-hf"):
    """Start the CoVe training process"""
    if not hf_token:
        return "❌ Please provide a HuggingFace token"
    
    # Set environment variables
    os.environ["HF_TOKEN"] = hf_token
    os.environ["SPACE_ID"] = "1"  # Indicates running in Spaces
    
    try:
        # Start training in a separate thread
        training_thread = threading.Thread(target=train_cove_model)
        training_thread.daemon = True
        training_thread.start()
        
        return "🚀 Training started! Check the logs below for progress."
    except Exception as e:
        return f"❌ Error starting training: {str(e)}"

def get_training_logs():
    """Get training logs"""
    log_file = "training.log"
    if os.path.exists(log_file):
        with open(log_file, 'r') as f:
            return f.read()
    return "No logs available yet. Training may not have started."

def test_model():
    """Test the trained model"""
    try:
        result = test_cove_inference()
        return result
    except Exception as e:
        return f"❌ Error testing model: {str(e)}"

# Create Gradio interface
with gr.Blocks(title="CoVe Fine-tuning on HuggingFace Spaces") as demo:
    gr.Markdown("# 🔗 Chain of Verification (CoVe) Fine-tuning")
    gr.Markdown("Fine-tune CodeLlama with Chain of Verification for better code reasoning")
    
    with gr.Tab("Setup & Training"):
        with gr.Row():
            with gr.Column():
                gr.Markdown("### 1. Setup")
                install_btn = gr.Button("Install Requirements", variant="secondary")
                install_output = gr.Textbox(label="Installation Status", lines=2)
                
                gr.Markdown("### 2. Training Configuration")
                hf_token = gr.Textbox(
                    label="HuggingFace Token (Write Access)", 
                    type="password",
                    placeholder="hf_xxxxxxxxxxxxxxxxxxxx"
                )
                model_name = gr.Textbox(
                    label="Base Model", 
                    value="codellama/CodeLlama-7b-Instruct-hf",
                    interactive=True
                )
                
                start_btn = gr.Button("Start Training", variant="primary")
                training_status = gr.Textbox(label="Training Status", lines=2)
            
            with gr.Column():
                gr.Markdown("### Training Progress")
                logs = gr.Textbox(
                    label="Training Logs", 
                    lines=15, 
                    max_lines=20,
                    value="Logs will appear here once training starts..."
                )
                refresh_logs_btn = gr.Button("Refresh Logs", variant="secondary")
    
    with gr.Tab("Model Testing"):
        gr.Markdown("### Test the Fine-tuned Model")
        test_btn = gr.Button("Test CoVe Model", variant="primary")
        test_output = gr.Textbox(label="Model Output", lines=10)
    
    with gr.Tab("Instructions"):
        gr.Markdown("""
        ## How to Use This Space
        
        ### Step 1: Get Your HuggingFace Token
        1. Go to [HuggingFace Settings](https://huggingface.co/settings/tokens)
        2. Create a new token with **Write** access
        3. Copy the token (starts with `hf_`)
        
        ### Step 2: Install Requirements
        - Click "Install Requirements" button
        - Wait for installation to complete
        
        ### Step 3: Start Training
        1. Paste your HuggingFace token
        2. Click "Start Training"
        3. Monitor progress in the logs
        
        ### Step 4: Test the Model
        - Once training is complete, test the model in the "Model Testing" tab
        
        ### Important Notes:
        - **GPU Required**: This requires a paid GPU Space (T4 minimum)
        - **Training Time**: Approximately 2-4 hours depending on GPU
        - **Storage**: The model will be saved to your HuggingFace account
        - **Monitoring**: Check logs regularly for progress updates
        
        ### Troubleshooting:
        - If training fails, check the logs for error messages
        - Ensure your HuggingFace token has write permissions
        - Make sure you have sufficient GPU memory
        """)
    
    # Event handlers
    install_btn.click(install_requirements, outputs=install_output)
    start_btn.click(start_training, inputs=[hf_token, model_name], outputs=training_status)
    refresh_logs_btn.click(get_training_logs, outputs=logs)
    test_btn.click(test_model, outputs=test_output)
    
    # Auto-refresh logs every 30 seconds during training
    # Instead of demo.load(..., every=30)
    gr.Timer(30, lambda: logs.update(value=get_training_logs()))


if __name__ == "__main__":
    demo.launch(share=False)