app_api_only.py

#!/usr/bin/env python3
"""
Textilindo AI Training API - Pure API Version
No HTML interfaces, only API endpoints for training and chat
"""

import os
import json
import logging
import torch
from pathlib import Path
from datetime import datetime
from typing import Dict, Any, Optional
from fastapi import FastAPI, HTTPException, BackgroundTasks
from pydantic import BaseModel
import uvicorn

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Initialize FastAPI app
app = FastAPI(
    title="Textilindo AI Training API",
    description="Pure API-based training system for Textilindo AI Assistant",
    version="1.0.0"
)

# Training status storage
training_status = {
    "is_training": False,
    "progress": 0,
    "status": "idle",
    "current_step": 0,
    "total_steps": 0,
    "loss": 0.0,
    "start_time": None,
    "end_time": None,
    "error": None
}

# Request/Response models
class TrainingRequest(BaseModel):
    model_name: str = "distilgpt2"
    dataset_path: str = "data/lora_dataset_20250910_145055.jsonl"
    config_path: str = "configs/training_config.yaml"
    max_samples: int = 20
    epochs: int = 1
    batch_size: int = 1
    learning_rate: float = 5e-5

class TrainingResponse(BaseModel):
    success: bool
    message: str
    training_id: str
    status: str

class ChatRequest(BaseModel):
    message: str
    conversation_id: Optional[str] = None

class ChatResponse(BaseModel):
    response: str
    conversation_id: str
    status: str = "success"

# API Information
@app.get("/")
async def api_info():
    """API information endpoint"""
    return {
        "name": "Textilindo AI Training API",
        "version": "1.0.0",
        "description": "Pure API-based training system for Textilindo AI Assistant",
        "hardware": "2 vCPU, 16 GB RAM (CPU basic)",
        "status": "ready",
        "endpoints": {
            "training": {
                "start": "POST /api/train/start",
                "status": "GET /api/train/status", 
                "data": "GET /api/train/data",
                "gpu": "GET /api/train/gpu",
                "test": "POST /api/train/test"
            },
            "chat": {
                "chat": "POST /chat",
                "health": "GET /health"
            }
        }
    }

# Health check
@app.get("/health")
async def health_check():
    """Health check endpoint"""
    return {
        "status": "healthy",
        "timestamp": datetime.now().isoformat(),
        "hardware": "2 vCPU, 16 GB RAM"
    }

# Training API endpoints
@app.post("/api/train/start", response_model=TrainingResponse)
async def start_training(request: TrainingRequest, background_tasks: BackgroundTasks):
    """Start training process"""
    global training_status
    
    if training_status["is_training"]:
        raise HTTPException(status_code=400, detail="Training already in progress")
    
    # Validate inputs
    if not Path(request.dataset_path).exists():
        raise HTTPException(status_code=404, detail=f"Dataset not found: {request.dataset_path}")
    
    if not Path(request.config_path).exists():
        raise HTTPException(status_code=404, detail=f"Config not found: {request.config_path}")
    
    # Start training in background
    training_id = f"train_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
    
    background_tasks.add_task(
        train_model_async,
        request.model_name,
        request.dataset_path,
        request.config_path,
        request.max_samples,
        request.epochs,
        request.batch_size,
        request.learning_rate
    )
    
    return TrainingResponse(
        success=True,
        message="Training started successfully",
        training_id=training_id,
        status="started"
    )

@app.get("/api/train/status")
async def get_training_status():
    """Get current training status"""
    return training_status

@app.get("/api/train/data")
async def get_training_data_info():
    """Get information about available training data"""
    data_dir = Path("data")
    if not data_dir.exists():
        return {"files": [], "count": 0}
    
    jsonl_files = list(data_dir.glob("*.jsonl"))
    files_info = []
    
    for file in jsonl_files:
        try:
            with open(file, 'r', encoding='utf-8') as f:
                lines = f.readlines()
            files_info.append({
                "name": file.name,
                "size": file.stat().st_size,
                "lines": len(lines)
            })
        except Exception as e:
            files_info.append({
                "name": file.name,
                "error": str(e)
            })
    
    return {
        "files": files_info,
        "count": len(jsonl_files)
    }

@app.get("/api/train/gpu")
async def get_gpu_info():
    """Get GPU information"""
    try:
        gpu_available = torch.cuda.is_available()
        if gpu_available:
            gpu_count = torch.cuda.device_count()
            gpu_name = torch.cuda.get_device_name(0)
            gpu_memory = torch.cuda.get_device_properties(0).total_memory / (1024**3)
            return {
                "available": True,
                "count": gpu_count,
                "name": gpu_name,
                "memory_gb": round(gpu_memory, 2)
            }
        else:
            return {"available": False}
    except Exception as e:
        return {"error": str(e)}

@app.post("/api/train/test")
async def test_trained_model():
    """Test the trained model"""
    model_path = "./models/textilindo-trained"
    if not Path(model_path).exists():
        return {"error": "No trained model found"}
    
    try:
        from transformers import AutoTokenizer, AutoModelForCausalLM
        
        tokenizer = AutoTokenizer.from_pretrained(model_path)
        model = AutoModelForCausalLM.from_pretrained(model_path)
        
        # Test prompt
        test_prompt = "Question: dimana lokasi textilindo? Answer:"
        inputs = tokenizer(test_prompt, return_tensors="pt")
        
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_length=inputs.input_ids.shape[1] + 30,
                temperature=0.7,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id
            )
        
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        return {
            "success": True,
            "test_prompt": test_prompt,
            "response": response,
            "model_path": model_path
        }
        
    except Exception as e:
        return {"error": str(e)}

# Chat API endpoint
@app.post("/chat", response_model=ChatResponse)
async def chat(request: ChatRequest):
    """Chat with the AI assistant"""
    try:
        # Simple mock response for now
        mock_responses = {
            "dimana lokasi textilindo": "Textilindo berkantor pusat di Jl. Raya Prancis No.39, Kosambi Tim., Kec. Kosambi, Kabupaten Tangerang, Banten 15213",
            "jam berapa textilindo beroperasional": "Jam operasional Senin-Jumat 08:00-17:00, Sabtu 08:00-12:00.",
            "berapa ketentuan pembelian": "Minimal order 1 roll per jenis kain",
            "apa ada gratis ongkir": "Gratis ongkir untuk order minimal 5 roll.",
            "apa bisa dikirimkan sample": "Hallo kak untuk sampel kita bisa kirimkan gratis ya kak 😊"
        }
        
        # Simple keyword matching
        user_lower = request.message.lower()
        response = "Halo! Saya adalah asisten AI Textilindo. Bagaimana saya bisa membantu Anda hari ini? 😊"
        
        for key, mock_response in mock_responses.items():
            if any(word in user_lower for word in key.split()):
                response = mock_response
                break
        
        return ChatResponse(
            response=response,
            conversation_id=request.conversation_id or "default",
            status="success"
        )
        
    except Exception as e:
        logger.error(f"Chat error: {e}")
        return ChatResponse(
            response="Maaf, terjadi kesalahan. Silakan coba lagi.",
            conversation_id=request.conversation_id or "default",
            status="error"
        )

# Training function
async def train_model_async(
    model_name: str,
    dataset_path: str,
    config_path: str,
    max_samples: int,
    epochs: int,
    batch_size: int,
    learning_rate: float
):
    """Async training function"""
    global training_status
    
    try:
        training_status.update({
            "is_training": True,
            "status": "starting",
            "progress": 0,
            "start_time": datetime.now().isoformat(),
            "error": None
        })
        
        logger.info("🚀 Starting training...")
        
        # Import training libraries
        from transformers import (
            AutoTokenizer, 
            AutoModelForCausalLM,
            TrainingArguments,
            Trainer,
            DataCollatorForLanguageModeling
        )
        from datasets import Dataset
        
        # Check GPU
        gpu_available = torch.cuda.is_available()
        logger.info(f"GPU available: {gpu_available}")
        
        # Load model and tokenizer
        logger.info(f"📥 Loading model: {model_name}")
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
        
        # Load model
        if gpu_available:
            model = AutoModelForCausalLM.from_pretrained(
                model_name,
                torch_dtype=torch.float16,
                device_map="auto"
            )
        else:
            model = AutoModelForCausalLM.from_pretrained(model_name)
        
        logger.info("✅ Model loaded successfully")
        
        # Load training data
        training_data = load_training_data(dataset_path, max_samples)
        if not training_data:
            raise Exception("No training data loaded")
        
        # Convert to dataset
        dataset = Dataset.from_list(training_data)
        
        def tokenize_function(examples):
            return tokenizer(
                examples["text"],
                truncation=True,
                padding=True,
                max_length=256,
                return_tensors="pt"
            )
        
        tokenized_dataset = dataset.map(tokenize_function, batched=True)
        
        # Training arguments
        training_args = TrainingArguments(
            output_dir="./models/textilindo-trained",
            num_train_epochs=epochs,
            per_device_train_batch_size=batch_size,
            gradient_accumulation_steps=2,
            learning_rate=learning_rate,
            warmup_steps=5,
            save_steps=10,
            logging_steps=1,
            save_total_limit=1,
            prediction_loss_only=True,
            remove_unused_columns=False,
            fp16=gpu_available,
            dataloader_pin_memory=gpu_available,
            report_to=None,
        )
        
        # Data collator
        data_collator = DataCollatorForLanguageModeling(
            tokenizer=tokenizer,
            mlm=False,
        )
        
        # Create trainer
        trainer = Trainer(
            model=model,
            args=training_args,
            train_dataset=tokenized_dataset,
            data_collator=data_collator,
            tokenizer=tokenizer,
        )
        
        # Start training
        training_status["status"] = "training"
        trainer.train()
        
        # Save model
        model.save_pretrained("./models/textilindo-trained")
        tokenizer.save_pretrained("./models/textilindo-trained")
        
        # Update status
        training_status.update({
            "is_training": False,
            "status": "completed",
            "progress": 100,
            "end_time": datetime.now().isoformat()
        })
        
        logger.info("✅ Training completed successfully!")
        
    except Exception as e:
        logger.error(f"Training failed: {e}")
        training_status.update({
            "is_training": False,
            "status": "failed",
            "error": str(e),
            "end_time": datetime.now().isoformat()
        })

def load_training_data(dataset_path: str, max_samples: int = 20) -> list:
    """Load training data from JSONL file"""
    data = []
    try:
        with open(dataset_path, 'r', encoding='utf-8') as f:
            for i, line in enumerate(f):
                if i >= max_samples:
                    break
                if line.strip():
                    item = json.loads(line)
                    # Create training text
                    instruction = item.get('instruction', '')
                    output = item.get('output', '')
                    text = f"Question: {instruction} Answer: {output}"
                    data.append({"text": text})
        logger.info(f"Loaded {len(data)} training samples")
        return data
    except Exception as e:
        logger.error(f"Error loading data: {e}")
        return []

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)