test_training_api.py

#!/usr/bin/env python3
"""
Test script for Textilindo AI Training API
Tests the training endpoints on Hugging Face Space
"""

import requests
import json
import time

# Base URL for the Hugging Face Space
BASE_URL = "https://harismlnaslm-Textilindo-AI.hf.space"

def test_gpu_status():
    """Test GPU availability"""
    print("🔍 Testing GPU Status...")
    try:
        response = requests.get(f"{BASE_URL}/api/train/gpu")
        if response.status_code == 200:
            data = response.json()
            print(f"✅ GPU Status: {data}")
            return data.get('available', False)
        else:
            print(f"❌ GPU Status Error: {response.status_code}")
            return False
    except Exception as e:
        print(f"❌ GPU Status Error: {e}")
        return False

def test_training_data():
    """Test training data availability"""
    print("\n📊 Testing Training Data...")
    try:
        response = requests.get(f"{BASE_URL}/api/train/data")
        if response.status_code == 200:
            data = response.json()
            print(f"✅ Training Data: {data['count']} files found")
            for file in data['files']:
                print(f"   📁 {file['name']}: {file['lines']} lines, {file['size']} bytes")
            return data
        else:
            print(f"❌ Training Data Error: {response.status_code}")
            return None
    except Exception as e:
        print(f"❌ Training Data Error: {e}")
        return None

def test_training_status():
    """Test current training status"""
    print("\n📈 Testing Training Status...")
    try:
        response = requests.get(f"{BASE_URL}/api/train/status")
        if response.status_code == 200:
            data = response.json()
            print(f"✅ Training Status: {data}")
            return data
        else:
            print(f"❌ Training Status Error: {response.status_code}")
            return None
    except Exception as e:
        print(f"❌ Training Status Error: {e}")
        return None

def start_training():
    """Start training with the largest dataset"""
    print("\n🚀 Starting Training...")
    
    # Use the largest dataset (lora_dataset_20250910_145055.jsonl)
    training_request = {
        "model_name": "distilgpt2",
        "dataset_path": "data/lora_dataset_20250910_145055.jsonl",
        "config_path": "configs/training_config.yaml",
        "max_samples": 20,  # Use more samples since we have 115 lines
        "epochs": 1,
        "batch_size": 1,
        "learning_rate": 5e-5
    }
    
    try:
        response = requests.post(
            f"{BASE_URL}/api/train/start",
            headers={"Content-Type": "application/json"},
            json=training_request
        )
        
        if response.status_code == 200:
            data = response.json()
            print(f"✅ Training Started: {data}")
            return True
        else:
            print(f"❌ Training Start Error: {response.status_code} - {response.text}")
            return False
    except Exception as e:
        print(f"❌ Training Start Error: {e}")
        return False

def monitor_training():
    """Monitor training progress"""
    print("\n⏱️ Monitoring Training Progress...")
    
    for i in range(10):  # Check 10 times
        try:
            response = requests.get(f"{BASE_URL}/api/train/status")
            if response.status_code == 200:
                data = response.json()
                print(f"📊 Status: {data['status']}, Progress: {data.get('progress', 0)}%, Loss: {data.get('loss', 0)}")
                
                if data['status'] in ['completed', 'failed', 'stopped']:
                    print(f"🏁 Training finished with status: {data['status']}")
                    return data
                
                time.sleep(5)  # Wait 5 seconds between checks
            else:
                print(f"❌ Status Check Error: {response.status_code}")
                break
        except Exception as e:
            print(f"❌ Status Check Error: {e}")
            break
    
    return None

def test_trained_model():
    """Test the trained model"""
    print("\n🧪 Testing Trained Model...")
    try:
        response = requests.post(f"{BASE_URL}/api/train/test")
        if response.status_code == 200:
            data = response.json()
            print(f"✅ Model Test: {data}")
            return data
        else:
            print(f"❌ Model Test Error: {response.status_code} - {response.text}")
            return None
    except Exception as e:
        print(f"❌ Model Test Error: {e}")
        return None

def main():
    """Main test function"""
    print("🤖 Textilindo AI Training API Test")
    print("=" * 50)
    
    # Test 1: GPU Status
    gpu_available = test_gpu_status()
    
    # Test 2: Training Data
    data_info = test_training_data()
    
    # Test 3: Current Status
    current_status = test_training_status()
    
    # Test 4: Start Training (if not already training)
    if current_status and not current_status.get('is_training', False):
        if start_training():
            # Test 5: Monitor Training
            final_status = monitor_training()
            
            # Test 6: Test Model (if training completed)
            if final_status and final_status.get('status') == 'completed':
                test_trained_model()
    else:
        print("\n⚠️ Training already in progress or failed to start")
    
    print("\n🎉 API Test Complete!")

if __name__ == "__main__":
    main()