app.py

import gradio as gr
import os
import json
import torch
from transformers import (
    AutoTokenizer, AutoModelForCausalLM, 
    TrainingArguments, Trainer, 
    DataCollatorForLanguageModeling,
    pipeline
)
from datasets import Dataset
from huggingface_hub import HfApi, login
import spaces
from typing import Optional, Dict, Any, List, Tuple
import logging
import traceback
from datetime import datetime
import random
import re
from faker import Faker
import hashlib
import time
from collections import defaultdict
from functools import wraps

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# ==================== RATE LIMITING ====================

class RateLimiter:
    """Token bucket rate limiter"""
    
    def __init__(self):
        self.requests = defaultdict(list)
        self.limits = {
            'synthetic_generation': {'calls': 10, 'period': 3600},
            'model_training': {'calls': 3, 'period': 3600},
            'model_inference': {'calls': 50, 'period': 3600},
        }
    
    def _get_user_id(self, request: gr.Request) -> str:
        if request:
            identifier = f"{request.client.host}_{request.headers.get('user-agent', '')}"
            return hashlib.md5(identifier.encode()).hexdigest()
        return "anonymous"
    
    def _clean_old_requests(self, user_id: str, endpoint: str):
        if user_id not in self.requests:
            return
        current_time = time.time()
        period = self.limits[endpoint]['period']
        self.requests[user_id] = [
            req for req in self.requests[user_id]
            if req['endpoint'] == endpoint and current_time - req['timestamp'] < period
        ]
    
    def check_rate_limit(self, user_id: str, endpoint: str) -> Tuple[bool, str]:
        self._clean_old_requests(user_id, endpoint)
        user_requests = [req for req in self.requests[user_id] if req['endpoint'] == endpoint]
        limit = self.limits[endpoint]['calls']
        period = self.limits[endpoint]['period']
        
        if len(user_requests) >= limit:
            time_until_reset = period - (time.time() - user_requests[0]['timestamp'])
            minutes = int(time_until_reset / 60)
            return False, f"⏱️ Rate limit exceeded! Please wait {minutes} minutes."
        
        self.requests[user_id].append({'endpoint': endpoint, 'timestamp': time.time()})
        remaining = limit - len(user_requests) - 1
        return True, f"✅ Request accepted ({remaining} remaining this hour)"

rate_limiter = RateLimiter()

def rate_limit(endpoint: str):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            request = kwargs.get('request', None)
            if request:
                user_id = rate_limiter._get_user_id(request)
                allowed, message = rate_limiter.check_rate_limit(user_id, endpoint)
                if not allowed:
                    return f"🚫 {message}"
            return func(*args, **kwargs)
        return wrapper
    return decorator

# ==================== AUTHENTICATION ====================

class AuthManager:
    def __init__(self):
        self.authenticated_tokens = {}
        self.token_expiry = 86400
    
    def validate_hf_token(self, token: str) -> Tuple[bool, str, Optional[str]]:
        try:
            if not token or not token.strip():
                return False, "❌ Please provide a HuggingFace token", None
            
            token_hash = hashlib.sha256(token.encode()).hexdigest()
            if token_hash in self.authenticated_tokens:
                cached = self.authenticated_tokens[token_hash]
                if time.time() - cached['timestamp'] < self.token_expiry:
                    return True, f"✅ Welcome back, {cached['username']}!", cached['username']
            
            api = HfApi(token=token)
            user_info = api.whoami()
            username = user_info.get('name', 'Anonymous Architect')
            
            self.authenticated_tokens[token_hash] = {
                'username': username,
                'timestamp': time.time()
            }
            
            return True, f"🎉 Welcome, {username}!", username
            
        except Exception as e:
            return False, f"🔐 Token validation failed: {str(e)}", None

auth_manager = AuthManager()

# ==================== ERROR HANDLING ====================

class ArchitechError(Exception):
    pass

class DataGenerationError(ArchitechError):
    pass

class ModelTrainingError(ArchitechError):
    pass

class ModelInferenceError(ArchitechError):
    pass

def handle_errors(error_type: str = "general"):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            try:
                return func(*args, **kwargs)
            except torch.cuda.OutOfMemoryError:
                return "🔥 **GPU Memory Overflow!** Try: smaller batch size, smaller model, or less data."
            except PermissionError:
                return "🔒 **Permission Denied!** Check your HuggingFace token has WRITE access."
            except ConnectionError:
                return "🌐 **Connection Issue!** Can't reach HuggingFace. Check your network."
            except ValueError as e:
                return f"⚠️ **Invalid Input!** {str(e)}"
            except (DataGenerationError, ModelTrainingError, ModelInferenceError) as e:
                return f"🔧 **Architech Error:** {str(e)}"
            except Exception as e:
                logger.error(f"Error in {func.__name__}: {traceback.format_exc()}")
                return f"💥 **Unexpected Error:** {str(e)}"
        return wrapper
    return decorator# ==================== SYNTHETIC DATA GENERATOR ====================

class SyntheticDataGenerator:
    def __init__(self):
        self.faker = Faker()
        self.generation_templates = {
            "conversational": [
                "Human: {question}\nAssistant: {answer}",
                "User: {question}\nBot: {answer}",
            ],
            "instruction": [
                "### Instruction:\n{instruction}\n\n### Response:\n{response}",
            ],
        }
        
        self.domain_knowledge = {
            "technology": {
                "topics": ["AI", "machine learning", "cloud computing"],
                "concepts": ["algorithms", "APIs", "databases"],
                "contexts": ["software development", "digital transformation"]
            },
            "healthcare": {
                "topics": ["telemedicine", "diagnostics", "patient care"],
                "concepts": ["treatments", "procedures"],
                "contexts": ["clinical practice", "patient education"]
            },
            "finance": {
                "topics": ["fintech", "investment", "risk management"],
                "concepts": ["portfolios", "compliance"],
                "contexts": ["financial advisory", "personal finance"]
            },
            "general": {
                "topics": ["communication", "problem-solving"],
                "concepts": ["strategies", "best practices"],
                "contexts": ["daily life", "personal growth"]
            }
        }
    
    def _generate_question(self, topic, concept, context):
        templates = [
            f"How does {concept} work in {context}?",
            f"What are the benefits of {concept} for {topic}?",
            f"Can you explain {concept}?",
            f"What's the best approach to {concept}?"
        ]
        return random.choice(templates)
    
    def _generate_answer(self, question, topic, concept):
        templates = [
            f"{concept} in {topic} works through strategic implementation. Key benefits include improved efficiency and better outcomes.",
            f"Great question! {concept} is fundamental because it addresses core challenges. Best practices include planning and testing.",
            f"When it comes to {concept}, consider scalability and performance. Success depends on proper implementation."
        ]
        return random.choice(templates)
    
    def _generate_single_example(self, task_desc, domain_data, templates, complexity):
        template = random.choice(templates)
        topic = random.choice(domain_data["topics"])
        concept = random.choice(domain_data["concepts"])
        context = random.choice(domain_data["contexts"])
        
        question = self._generate_question(topic, concept, context)
        answer = self._generate_answer(question, topic, concept)
        
        text = template.format(question=question, answer=answer)
        return {"text": text}
    
    @handle_errors("data_generation")
    def generate_synthetic_dataset(
        self,
        task_description: str,
        domain: str,
        dataset_size: int = 100,
        format_type: str = "conversational",
        complexity: str = "medium",
        progress=gr.Progress()
    ) -> str:
        if not task_description or len(task_description.strip()) < 10:
            raise DataGenerationError("Task description too short! Need at least 10 characters.")
        
        if dataset_size < 10 or dataset_size > 1000:
            raise DataGenerationError("Dataset size must be between 10 and 1000.")
        
        progress(0.1, f"🎯 Generating {dataset_size} examples...")
        
        domain_data = self.domain_knowledge.get(domain, self.domain_knowledge["general"])
        templates = self.generation_templates.get(format_type, self.generation_templates["conversational"])
        
        synthetic_data = []
        for i in range(dataset_size):
            if i % 20 == 0:
                progress(0.1 + (0.7 * i / dataset_size), f"📝 Creating {i+1}/{dataset_size}...")
            
            example = self._generate_single_example(task_description, domain_data, templates, complexity)
            synthetic_data.append(example)
        
        os.makedirs("./synthetic_datasets", exist_ok=True)
        dataset_filename = f"synthetic_{domain}_{format_type}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
        dataset_path = os.path.join("./synthetic_datasets", dataset_filename)
        
        with open(dataset_path, 'w') as f:
            json.dump(synthetic_data, f, indent=2)
        
        preview = "\n\n---\n\n".join([ex["text"] for ex in synthetic_data[:3]])
        
        return f"""🎊 **SYNTHETIC DATASET GENERATED!**

**Dataset Details:**
- 📊 Size: {len(synthetic_data)} examples
- 🎯 Domain: {domain.title()}
- 📝 Format: {format_type.title()}
- 💾 Saved as: `{dataset_filename}`

**Preview (First 3 Examples):**

{preview}

**Next Steps:** Use this in the 'Train Model' or 'Test Model' tabs!"""# ==================== MODEL INFERENCE ====================

class ModelInference:
    def __init__(self):
        self.loaded_models = {}
    
    @handle_errors("inference")
    def load_model(self, model_name: str, hf_token: str, progress=gr.Progress()) -> str:
        progress(0.1, "🔍 Locating your model...")
        
        is_valid, message, username = auth_manager.validate_hf_token(hf_token)
        if not is_valid:
            raise ModelInferenceError(message)
        
        full_model_name = f"{username}/{model_name}" if "/" not in model_name else model_name
        
        progress(0.3, "📥 Downloading model...")
        
        try:
            tokenizer = AutoTokenizer.from_pretrained(full_model_name, token=hf_token)
            model = AutoModelForCausalLM.from_pretrained(
                full_model_name,
                token=hf_token,
                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
                device_map="auto" if torch.cuda.is_available() else None
            )
            
            self.loaded_models[model_name] = {
                'model': model,
                'tokenizer': tokenizer,
                'pipeline': pipeline('text-generation', model=model, tokenizer=tokenizer)
            }
            
            progress(1.0, "✅ Model loaded!")
            return f"✅ **Model Loaded Successfully!**\n\nModel: `{full_model_name}`\n\nReady for inference!"
            
        except Exception as e:
            raise ModelInferenceError(f"Failed to load model: {str(e)}")
    
    @handle_errors("inference")
    def generate_text(
        self,
        model_name: str,
        prompt: str,
        max_length: int = 100,
        temperature: float = 0.7,
        top_p: float = 0.9
    ) -> str:
        if model_name not in self.loaded_models:
            raise ModelInferenceError("Model not loaded! Please load the model first.")
        
        if not prompt or len(prompt.strip()) < 3:
            raise ModelInferenceError("Prompt too short! Please provide at least 3 characters.")
        
        pipe = self.loaded_models[model_name]['pipeline']
        
        result = pipe(
            prompt,
            max_length=max_length,
            temperature=temperature,
            top_p=top_p,
            do_sample=True,
            num_return_sequences=1
        )
        
        generated_text = result[0]['generated_text']
        
        return f"""**🎯 Generated Response:**

{generated_text}

---
*Model: {model_name} | Length: {len(generated_text)} chars*"""

model_inference = ModelInference()# ==================== ARCHITECH AGENT ====================

class ArchitechAgent:
    def __init__(self):
        self.hf_api = HfApi()
        self.synthetic_generator = SyntheticDataGenerator()
        self.personality_responses = [
            "🎯 Let's cook up some AI magic!",
            "🚀 Time to turn your vision into reality!",
            "🧠 Let's architect some brilliance!",
        ]
    
    def get_personality_response(self) -> str:
        return random.choice(self.personality_responses)
    
    @rate_limit('synthetic_generation')
    @handle_errors("data_generation")
    def generate_synthetic_dataset_wrapper(self, *args, **kwargs):
        return self.synthetic_generator.generate_synthetic_dataset(*args, **kwargs)
    
    @spaces.GPU
    @rate_limit('model_training')
    @handle_errors("training")
    def train_custom_model(
        self,
        task_description: str,
        training_data: str,
        model_name: str,
        hf_token: str,
        base_model: str = "distilgpt2",
        use_synthetic_data: bool = True,
        synthetic_domain: str = "general",
        synthetic_size: int = 100,
        learning_rate: float = 2e-4,
        num_epochs: int = 3,
        batch_size: int = 2,
        progress=gr.Progress()
    ) -> str:
        
        is_valid, message, username = auth_manager.validate_hf_token(hf_token)
        if not is_valid:
            raise ModelTrainingError(message)
        
        progress(0.1, "🧠 Loading base model...")
        
        tokenizer = AutoTokenizer.from_pretrained(base_model)
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
        
        model = AutoModelForCausalLM.from_pretrained(
            base_model,
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
            device_map="auto" if torch.cuda.is_available() else None
        )
        
        if use_synthetic_data:
            progress(0.2, "🎨 Generating synthetic data...")
            result = self.synthetic_generator.generate_synthetic_dataset(
                task_description, synthetic_domain, synthetic_size, "conversational", "medium", progress
            )
            
            dataset_files = [f for f in os.listdir("./synthetic_datasets") if f.endswith('.json')]
            if not dataset_files:
                raise ModelTrainingError("No synthetic dataset found!")
            
            latest_dataset = max(dataset_files, key=lambda x: os.path.getctime(os.path.join("./synthetic_datasets", x)))
            with open(os.path.join("./synthetic_datasets", latest_dataset), 'r') as f:
                synthetic_data = json.load(f)
            texts = [item["text"] for item in synthetic_data]
        else:
            # Check if training_data is a file path or raw text
            if training_data.strip().endswith('.json') and os.path.exists(training_data.strip()):
                # Load from file
                texts = dataset_manager.load_dataset_for_training(training_data.strip())
            else:
                # Parse as raw text
                texts = [t.strip() for t in training_data.split("\n\n") if t.strip()]
        
        if not texts:
            raise ModelTrainingError("No training data available!")
        
        progress(0.3, f"✨ Tokenizing {len(texts)} examples...")
        
        dataset = Dataset.from_dict({"text": texts})
        
        def tokenize_function(examples):
            return tokenizer(examples["text"], truncation=True, padding=True, max_length=256)
        
        tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=["text"])
        
        progress(0.4, "⚙️ Configuring training...")
        
        training_args = TrainingArguments(
            output_dir=f"./results_{model_name}",
            num_train_epochs=num_epochs,
            per_device_train_batch_size=batch_size,
            gradient_accumulation_steps=4,
            learning_rate=learning_rate,
            logging_steps=50,
            save_steps=500,
            save_total_limit=2,
            fp16=torch.cuda.is_available(),
            report_to="none"
        )
        
        data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
        
        trainer = Trainer(
            model=model,
            args=training_args,
            train_dataset=tokenized_dataset,
            data_collator=data_collator,
        )
        
        progress(0.6, "💪 Training in progress...")
        trainer.train()
        
        progress(0.8, "💾 Saving model...")
        output_dir = f"./trained_{model_name}"
        trainer.save_model(output_dir)
        tokenizer.save_pretrained(output_dir)
        
        progress(0.9, "📤 Pushing to HuggingFace...")
        
        try:
            login(token=hf_token)
            
            # Try uploading with retries
            max_retries = 3
            for attempt in range(max_retries):
                try:
                    progress(0.9 + (attempt * 0.03), f"📤 Upload attempt {attempt + 1}/{max_retries}...")
                    
                    # Push model with timeout
                    model.push_to_hub(
                        model_name, 
                        token=hf_token,
                        max_shard_size="500MB",
                        safe_serialization=True
                    )
                    tokenizer.push_to_hub(model_name, token=hf_token)
                    
                    hub_url = f"https://huggingface.co/{username}/{model_name}"
                    
                    return f"""🎉 **TRAINING COMPLETE!**

✅ Training successful
💾 Model saved locally
📤 Pushed to Hub
🔗 **Your model:** {hub_url}

**Stats:**
- Examples: {len(texts)}
- Epochs: {num_epochs}
- Learning rate: {learning_rate}

**Test it in the 'Test Model' tab!**"""
                
                except Exception as upload_error:
                    if attempt < max_retries - 1:
                        logger.warning(f"Upload attempt {attempt + 1} failed: {upload_error}")
                        time.sleep(5)  # Wait before retry
                        continue
                    else:
                        raise upload_error
        
        except Exception as e:
            logger.error(f"Upload failed after retries: {e}")
            
            # Provide manual upload instructions
            return f"""✅ **TRAINING COMPLETE!** (Upload timed out)

💾 Model saved locally at: `{output_dir}`

**Manual Upload Instructions:**
1. Download your Space's files (or access via SSH if enabled)
2. Run this command locally:
```bash
huggingface-cli upload {username}/{model_name} {output_dir}
```

Or use the Python API:
```python
from huggingface_hub import HfApi
api = HfApi()
api.upload_folder(
    folder_path="{output_dir}",
    repo_id="{username}/{model_name}",
    token="YOUR_TOKEN"
)
```

**Stats:**
- Examples: {len(texts)}
- Epochs: {num_epochs}
- Model saved successfully!

**You can still test it locally or manually upload!**"""# ==================== MODEL MANAGEMENT ====================

import zipfile
import shutil
from pathlib import Path

class ModelManager:
    def __init__(self):
        self.models_dir = Path("./saved_models")
        self.models_dir.mkdir(exist_ok=True)
    
    @handle_errors("model_management")
    def create_model_zip(self, model_path: str, model_name: str) -> Tuple[str, str]:
        """Create a downloadable zip of a trained model"""
        if not os.path.exists(model_path):
            raise ArchitechError(f"Model path not found: {model_path}")
        
        zip_filename = f"{model_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.zip"
        zip_path = os.path.join(self.models_dir, zip_filename)
        
        # Create zip file
        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
            for root, dirs, files in os.walk(model_path):
                for file in files:
                    file_path = os.path.join(root, file)
                    arcname = os.path.relpath(file_path, model_path)
                    zipf.write(file_path, arcname)
        
        file_size = os.path.getsize(zip_path) / (1024 * 1024)  # MB
        
        return zip_path, f"✅ Created {zip_filename} ({file_size:.2f} MB)"
    
    @handle_errors("model_management")
    def extract_model_zip(self, zip_file, progress=gr.Progress()) -> str:
        """Extract uploaded model zip"""
        if zip_file is None:
            raise ArchitechError("No file uploaded!")
        
        progress(0.1, "📦 Extracting model archive...")
        
        # Get filename
        zip_filename = Path(zip_file.name).name
        model_name = zip_filename.replace('.zip', '')
        extract_path = os.path.join("./uploaded_models", model_name)
        
        os.makedirs(extract_path, exist_ok=True)
        
        progress(0.3, "📂 Unpacking files...")
        
        # Extract zip
        with zipfile.ZipFile(zip_file.name, 'r') as zip_ref:
            zip_ref.extractall(extract_path)
        
        progress(0.7, "🔍 Validating model files...")
        
        # Check for required files
        files = os.listdir(extract_path)
        has_model = any('pytorch_model' in f or 'model.safetensors' in f for f in files)
        has_config = 'config.json' in files
        has_tokenizer = any('tokenizer' in f for f in files)
        
        validation_status = []
        if has_model:
            validation_status.append("✅ Model weights found")
        else:
            validation_status.append("⚠️ Model weights not found")
        
        if has_config:
            validation_status.append("✅ Config file found")
        else:
            validation_status.append("⚠️ Config file not found")
        
        if has_tokenizer:
            validation_status.append("✅ Tokenizer found")
        else:
            validation_status.append("⚠️ Tokenizer not found")
        
        progress(1.0, "✅ Extraction complete!")
        
        return f"""🎉 **Model Uploaded Successfully!**

**Extracted to:** `{extract_path}`

**Validation:**
{chr(10).join(validation_status)}

**Files found:** {len(files)} files

**You can now:**
1. Use this model in the Test Model tab
2. Continue training from this checkpoint
3. Push to HuggingFace Hub

*Model path: `{extract_path}`*"""
    
    def list_local_models(self) -> str:
        """List all locally saved models"""
        trained_models = []
        uploaded_models = []
        
        # Check trained models
        if os.path.exists("./"):
            for item in os.listdir("./"):
                if item.startswith("trained_") and os.path.isdir(item):
                    size = sum(
                        os.path.getsize(os.path.join(dirpath, filename))
                        for dirpath, dirnames, filenames in os.walk(item)
                        for filename in filenames
                    ) / (1024 * 1024)
                    trained_models.append(f"- `{item}` ({size:.2f} MB)")
        
        # Check uploaded models
        if os.path.exists("./uploaded_models"):
            for item in os.listdir("./uploaded_models"):
                path = os.path.join("./uploaded_models", item)
                if os.path.isdir(path):
                    size = sum(
                        os.path.getsize(os.path.join(dirpath, filename))
                        for dirpath, dirnames, filenames in os.walk(path)
                        for filename in filenames
                    ) / (1024 * 1024)
                    uploaded_models.append(f"- `{item}` ({size:.2f} MB)")
        
        result = "## 📦 Local Models\n\n"
        
        if trained_models:
            result += "### Trained Models:\n" + "\n".join(trained_models) + "\n\n"
        else:
            result += "### Trained Models:\n*No trained models found*\n\n"
        
        if uploaded_models:
            result += "### Uploaded Models:\n" + "\n".join(uploaded_models) + "\n\n"
        else:
            result += "### Uploaded Models:\n*No uploaded models found*\n\n"
        
        return result
    
    @handle_errors("model_management")
    def delete_model(self, model_path: str) -> str:
        """Delete a local model"""
        if not os.path.exists(model_path):
            raise ArchitechError(f"Model not found: {model_path}")
        
        shutil.rmtree(model_path)
        return f"✅ Deleted: {model_path}"

model_manager = ModelManager()

# Add this to the Gradio interface creation function
# Insert this tab after the "Test Model" tab and before "About"

def add_model_management_tab():
    """Add Model Management tab to Gradio interface"""
    with gr.Tab("💾 Model Management"):
        gr.Markdown("""
        ### Manage Your Models
        Upload, download, and organize your trained models
        """)
        
        with gr.Row():
            # Upload Section
            with gr.Column():
                gr.Markdown("### 📤 Upload Model")
                
                upload_file = gr.File(
                    label="Upload Model ZIP",
                    file_types=[".zip"],
                    type="filepath"
                )
                
                upload_btn = gr.Button("📦 Extract and Save", variant="primary")
                upload_output = gr.Markdown()
                
                upload_btn.click(
                    fn=model_manager.extract_model_zip,
                    inputs=[upload_file],
                    outputs=upload_output
                )
            
            # Download Section
            with gr.Column():
                gr.Markdown("### 📥 Download Model")
                
                model_path_input = gr.Textbox(
                    label="Model Path",
                    placeholder="e.g., ./trained_my-model or ./uploaded_models/my-model",
                    info="Path to the model directory you want to download"
                )
                
                model_name_input = gr.Textbox(
                    label="Archive Name",
                    placeholder="e.g., my-awesome-model",
                    info="Name for the downloaded zip file"
                )
                
                download_btn = gr.Button("📦 Create ZIP", variant="primary")
                download_file = gr.File(label="Download")
                download_output = gr.Markdown()
                
                def create_and_return_zip(model_path, model_name):
                    zip_path, message = model_manager.create_model_zip(model_path, model_name)
                    return zip_path, message
                
                download_btn.click(
                    fn=create_and_return_zip,
                    inputs=[model_path_input, model_name_input],
                    outputs=[download_file, download_output]
                )
        
        gr.Markdown("---")
        
        # List and Delete Section
        with gr.Row():
            with gr.Column():
                gr.Markdown("### 📋 Your Models")
                
                refresh_btn = gr.Button("🔄 Refresh List", variant="secondary")
                models_list = gr.Markdown()
                
                refresh_btn.click(
                    fn=model_manager.list_local_models,
                    inputs=[],
                    outputs=models_list
                )
                
                # Auto-load on tab open
                models_list.value = model_manager.list_local_models()
            
            with gr.Column():
                gr.Markdown("### 🗑️ Delete Model")
                
                delete_path = gr.Textbox(
                    label="Model Path to Delete",
                    placeholder="e.g., ./trained_my-model"
                )
                
                delete_btn = gr.Button("🗑️ Delete Model", variant="stop")
                delete_output = gr.Markdown()
                
                delete_btn.click(
                    fn=model_manager.delete_model,
                    inputs=[delete_path],
                    outputs=delete_output
                )
        
        gr.Markdown("""
        ---
        ### 💡 Tips:
        - **Upload:** Upload model zips from other systems or backups
        - **Download:** Create portable archives of your trained models
        - **Organize:** Keep your workspace tidy by managing local models
        - **Backup:** Download important models before deleting them
        
        *Note: Uploaded/downloaded models persist only during your session unless you have persistent storage configured.*
        """)

# This function should be called in create_gradio_interface()
# Add it right before the "About" tab# ==================== DATASET MANAGER ====================

class DatasetManager:
    def __init__(self):
        self.datasets_dir = Path("./synthetic_datasets")
        self.datasets_dir.mkdir(exist_ok=True)
    
    def list_available_datasets(self) -> List[Tuple[str, str]]:
        """List all available synthetic datasets"""
        datasets = []
        if self.datasets_dir.exists():
            for file in self.datasets_dir.glob("*.json"):
                datasets.append((file.name, str(file)))
        return datasets
    
    def get_dataset_preview(self, dataset_path: str) -> str:
        """Get preview of dataset contents"""
        try:
            with open(dataset_path, 'r') as f:
                data = json.load(f)
            
            if not data:
                return "Dataset is empty"
            
            preview = f"**Dataset:** `{Path(dataset_path).name}`\n\n"
            preview += f"**Total Examples:** {len(data)}\n\n"
            preview += "**First 3 Examples:**\n\n"
            
            for i, example in enumerate(data[:3], 1):
                preview += f"**Example {i}:**\n```\n{example.get('text', 'No text field')}\n```\n\n"
            
            return preview
        except Exception as e:
            return f"Error loading dataset: {str(e)}"
    
    def load_dataset_for_training(self, dataset_path: str) -> List[str]:
        """Load dataset texts for training"""
        with open(dataset_path, 'r') as f:
            data = json.load(f)
        return [item["text"] for item in data if "text" in item]

dataset_manager = DatasetManager()

# ==================== REPOSITORY CHAT SYSTEM ====================

class RepositoryChat:
    def __init__(self):
        self.hf_api = HfApi()
        self.chat_history = []
        self.current_user_token = None
        self.current_username = None
    
    def initialize_session(self, hf_token: str) -> Tuple[bool, str]:
        """Initialize chat session with HF token"""
        is_valid, message, username = auth_manager.validate_hf_token(hf_token)
        if is_valid:
            self.current_user_token = hf_token
            self.current_username = username
            self.chat_history = []
        return is_valid, message
    
    @handle_errors("repository_chat")
    def list_user_models(self) -> str:
        """List all models in user's HuggingFace account"""
        if not self.current_user_token:
            raise ArchitechError("Please initialize session with your HuggingFace token first!")
        
        try:
            models = self.hf_api.list_models(author=self.current_username, token=self.current_user_token)
            model_list = list(models)
            
            if not model_list:
                return f"📭 No models found in {self.current_username}'s account"
            
            result = f"## 🤖 Your Models ({len(model_list)})\n\n"
            
            for model in model_list[:20]:  # Limit to 20 for display
                model_id = model.modelId
                downloads = getattr(model, 'downloads', 0)
                likes = getattr(model, 'likes', 0)
                result += f"- **{model_id}**\n"
                result += f"  - Downloads: {downloads} | Likes: {likes}\n"
                result += f"  - [View on Hub](https://huggingface.co/{model_id})\n\n"
            
            return result
        except Exception as e:
            return f"Error fetching models: {str(e)}"
    
    @handle_errors("repository_chat")
    def list_user_datasets(self) -> str:
        """List all datasets in user's HuggingFace account"""
        if not self.current_user_token:
            raise ArchitechError("Please initialize session first!")
        
        try:
            datasets = self.hf_api.list_datasets(author=self.current_username, token=self.current_user_token)
            dataset_list = list(datasets)
            
            if not dataset_list:
                return f"📭 No datasets found in {self.current_username}'s account"
            
            result = f"## 📊 Your Datasets ({len(dataset_list)})\n\n"
            
            for dataset in dataset_list[:20]:
                dataset_id = dataset.id
                downloads = getattr(dataset, 'downloads', 0)
                result += f"- **{dataset_id}**\n"
                result += f"  - Downloads: {downloads}\n"
                result += f"  - [View on Hub](https://huggingface.co/datasets/{dataset_id})\n\n"
            
            return result
        except Exception as e:
            return f"Error fetching datasets: {str(e)}"
    
    @handle_errors("repository_chat")
    def get_model_info(self, model_id: str) -> str:
        """Get detailed information about a specific model"""
        if not self.current_user_token:
            raise ArchitechError("Please initialize session first!")
        
        try:
            # Add username if not in model_id
            if "/" not in model_id and self.current_username:
                model_id = f"{self.current_username}/{model_id}"
            
            model_info = self.hf_api.model_info(model_id, token=self.current_user_token)
            
            result = f"## 🤖 Model: {model_id}\n\n"
            result += f"**Model ID:** {model_info.modelId}\n"
            result += f"**Downloads:** {getattr(model_info, 'downloads', 0)}\n"
            result += f"**Likes:** {getattr(model_info, 'likes', 0)}\n"
            result += f"**Created:** {getattr(model_info, 'created_at', 'Unknown')}\n"
            result += f"**Last Modified:** {getattr(model_info, 'last_modified', 'Unknown')}\n\n"
            
            if hasattr(model_info, 'tags') and model_info.tags:
                result += f"**Tags:** {', '.join(model_info.tags[:10])}\n\n"
            
            result += f"**🔗 [View on HuggingFace](https://huggingface.co/{model_id})**\n"
            
            return result
        except Exception as e:
            return f"Error fetching model info: {str(e)}"
    
    @handle_errors("repository_chat")
    def delete_repo(self, repo_id: str, repo_type: str = "model") -> str:
        """Delete a repository (model or dataset)"""
        if not self.current_user_token:
            raise ArchitechError("Please initialize session first!")
        
        # Add username if not in repo_id
        if "/" not in repo_id and self.current_username:
            repo_id = f"{self.current_username}/{repo_id}"
        
        try:
            self.hf_api.delete_repo(
                repo_id=repo_id,
                token=self.current_user_token,
                repo_type=repo_type
            )
            return f"✅ Successfully deleted {repo_type}: {repo_id}"
        except Exception as e:
            return f"❌ Error deleting {repo_type}: {str(e)}"
    
    @handle_errors("repository_chat")
    def chat_with_repos(self, user_message: str) -> str:
        """Conversational interface for repository management"""
        if not self.current_user_token:
            return "⚠️ Please initialize your session with a HuggingFace token first!"
        
        # Add to history
        self.chat_history.append({"role": "user", "content": user_message})
        
        # Parse intent
        message_lower = user_message.lower()
        
        response = ""
        
        # List models
        if any(word in message_lower for word in ["list models", "show models", "my models", "what models"]):
            response = self.list_user_models()
        
        # List datasets
        elif any(word in message_lower for word in ["list datasets", "show datasets", "my datasets", "what datasets"]):
            response = self.list_user_datasets()
        
        # Model info
        elif any(word in message_lower for word in ["info about", "details about", "tell me about", "information on"]):
            # Extract model name (simple extraction)
            words = user_message.split()
            if len(words) > 2:
                potential_model = words[-1].strip("?.,!")
                response = self.get_model_info(potential_model)
            else:
                response = "Please specify which model you want info about. Example: 'info about my-model-name'"
        
        # Delete model
        elif "delete" in message_lower and "model" in message_lower:
            words = user_message.split()
            if len(words) > 2:
                model_name = words[-1].strip("?.,!")
                response = f"⚠️ Are you sure you want to delete model '{model_name}'? This action cannot be undone!\n\n"
                response += "To confirm, use the Delete Repository section below."
            else:
                response = "Please specify which model to delete. Example: 'delete model my-model-name'"
        
        # General help
        elif any(word in message_lower for word in ["help", "what can you do", "commands"]):
            response = """## 🤖 Architech Repository Assistant

I can help you manage your HuggingFace repositories! Here's what I can do:

**📋 Listing:**
- "List my models" - Show all your models
- "Show my datasets" - Show all your datasets

**ℹ️ Information:**
- "Info about [model-name]" - Get details about a specific model
- "Tell me about [model-name]" - Model statistics and info

**🗑️ Management:**
- Use the Delete Repository section to remove models/datasets

**💡 Tips:**
- I have access to your HuggingFace account
- I can see all your public and private repos
- All actions respect your permissions

Try asking: "List my models" or "Show my datasets"!"""
        
        # Default response
        else:
            response = f"""I'm not sure what you want to do. 

**Quick Commands:**
- "List my models"
- "Show my datasets"
- "Info about [model-name]"
- "Help" for full command list

What would you like to do?"""
        
        # Add to history
        self.chat_history.append({"role": "assistant", "content": response})
        
        return response
    
    def get_chat_history_display(self) -> List[Tuple[str, str]]:
        """Format chat history for Gradio ChatBot"""
        history = []
        for i in range(0, len(self.chat_history), 2):
            if i + 1 < len(self.chat_history):
                user_msg = self.chat_history[i]["content"]
                bot_msg = self.chat_history[i + 1]["content"]
                history.append((user_msg, bot_msg))
        return history

repo_chat = RepositoryChat()# # ==================== MODEL CARD & PAPER GENERATOR ====================

class DocumentationGenerator:
    def __init__(self):
        self.templates_dir = Path("./generated_docs")
        self.templates_dir.mkdir(exist_ok=True)
    
    def generate_model_card(
        self,
        model_name: str,
        task_description: str,
        base_model: str,
        dataset_size: int,
        training_params: Dict[str, Any],
        domain: str = "general",
        intended_use: str = "",
        limitations: str = "",
        ethical_considerations: str = ""
    ) -> str:
        """Generate a comprehensive model card following HuggingFace standards"""
        
        timestamp = datetime.now().strftime("%Y-%m-%d")
        
        model_card = f"""---
language: en
license: mit
tags:
- text-generation
- custom-model
- architech
- {domain}
datasets:
- synthetic-data
metrics:
- perplexity
model-index:
- name: {model_name}
  results: []
---

# {model_name}

## Model Description

**{model_name}** is a fine-tuned language model created using Architech AI Model Architect.

### Model Details

- **Developed by:** Architech User
- **Model type:** Causal Language Model
- **Language(s):** English
- **Base Model:** {base_model}
- **License:** MIT
- **Finetuned from:** {base_model}

### Model Purpose

{task_description}

## Training Details

### Training Data

This model was trained on a synthetic dataset specifically generated for this task:

- **Dataset Size:** {dataset_size} examples
- **Domain:** {domain.title()}
- **Data Generation:** Architech Synthetic Data Generator
- **Data Format:** Conversational pairs / Instruction-response format

The training data was synthetically generated to ensure:
- Domain-specific vocabulary and concepts
- Natural language variations
- Task-relevant examples
- Ethical and unbiased content

### Training Procedure

**Training Hyperparameters:**

- **Base Model:** {base_model}
- **Training Examples:** {dataset_size}
- **Epochs:** {training_params.get('epochs', 'N/A')}
- **Learning Rate:** {training_params.get('learning_rate', 'N/A')}
- **Batch Size:** {training_params.get('batch_size', 'N/A')}
- **Gradient Accumulation Steps:** {training_params.get('gradient_accumulation', 4)}
- **Optimizer:** AdamW
- **Training Precision:** FP16 (if GPU available)

**Training Infrastructure:**

- **Framework:** HuggingFace Transformers
- **Training Tool:** Architech AI Model Architect
- **Hardware:** {training_params.get('hardware', 'GPU/CPU auto-detected')}

## Intended Use

### Direct Use

{intended_use if intended_use else f'''This model is designed for {task_description.lower()}. It can be used directly for:

- Text generation in the {domain} domain
- Conversational AI applications
- Task-specific completion and assistance
- Research and experimentation'''}

### Downstream Use

This model can be further fine-tuned for:
- More specialized tasks within the {domain} domain
- Multi-turn conversations
- Domain-specific applications

### Out-of-Scope Use

This model should NOT be used for:
- Medical, legal, or financial advice without human oversight
- Safety-critical applications
- Decision-making without human review
- Generating harmful, biased, or unethical content

## Bias, Risks, and Limitations

{limitations if limitations else f'''### Known Limitations

- Trained on synthetic data, which may not capture all real-world nuances
- Limited to {dataset_size} training examples
- May produce inconsistent outputs on topics outside training domain
- Should not be considered a source of factual information without verification

### Recommendations

Users should:
- Validate outputs for accuracy and appropriateness
- Not rely solely on this model for critical decisions
- Be aware of potential biases in generated content
- Use human oversight for production applications'''}

## Ethical Considerations

{ethical_considerations if ethical_considerations else '''This model was developed with ethical AI principles in mind:

- Training data was synthetically generated to avoid privacy issues
- No personally identifiable information was used in training
- Content generation should be monitored for potential misuse
- Users are responsible for ensuring ethical use of generated content'''}

## How to Use

### Loading the Model

```python
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("{model_name}")
model = AutoModelForCausalLM.from_pretrained("{model_name}")

# Generate text
inputs = tokenizer("Your prompt here", return_tensors="pt")
outputs = model.generate(**inputs, max_length=100)
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(generated_text)
```

### Using with Pipeline

```python
from transformers import pipeline

generator = pipeline('text-generation', model='{model_name}')
result = generator("Your prompt here", max_length=100)
print(result[0]['generated_text'])
```

## Model Performance

Performance metrics will vary based on specific use case and evaluation criteria.

### Training Loss

Training completed successfully with the model converging appropriately for the given dataset size and complexity.

## Environmental Impact

- **Training Time:** Approximately {training_params.get('training_time', 'varies')} minutes
- **Hardware:** {training_params.get('hardware', 'GPU/CPU')}
- **Carbon Emissions:** Minimal due to efficient training approach

## Technical Specifications

### Model Architecture

Based on {base_model} architecture with task-specific fine-tuning.

### Compute Infrastructure

- **Training Platform:** HuggingFace Spaces / Architech
- **Framework:** PyTorch + Transformers
- **Optimization:** Gradient accumulation for memory efficiency

## Citation

If you use this model, please cite:

```bibtex
@misc{{{model_name.replace('-', '_')},
  author = {{Architech User}},
  title = {{{model_name}}},
  year = {{{datetime.now().year}}},
  publisher = {{HuggingFace}},
  howpublished = {{\\url{{https://huggingface.co/your-username/{model_name}}}}}
}}
```

## Model Card Authors

- Generated by: Architech AI Model Architect
- Date: {timestamp}

## Model Card Contact

For questions or feedback about this model, please open an issue in the model repository.

---

*This model card was automatically generated by Architech AI Model Architect. Please review and customize as needed.*
"""
        
        # Save model card
        card_path = self.templates_dir / f"{model_name}_model_card.md"
        with open(card_path, 'w') as f:
            f.write(model_card)
        
        return model_card, str(card_path)
    
    def generate_research_paper(
        self,
        model_name: str,
        task_description: str,
        base_model: str,
        dataset_size: int,
        training_params: Dict[str, Any],
        domain: str = "general",
        methodology_notes: str = "",
        results_summary: str = ""
    ) -> str:
        """Generate a research paper documenting the model"""
        
        timestamp = datetime.now().strftime("%B %Y")
        
        paper = f"""# Fine-Tuning {base_model} for {task_description}: A Synthetic Data Approach

**Authors:** Architech User  
**Date:** {timestamp}  
**Model:** {model_name}

---

## Abstract

We present **{model_name}**, a fine-tuned language model specifically designed for {task_description.lower()}. 
This work demonstrates the effectiveness of synthetic data generation for domain-specific language model adaptation. 
Using {dataset_size} synthetically generated examples, we fine-tuned {base_model} to create a specialized model 
for the {domain} domain. Our approach leverages automated data generation techniques to overcome the common challenge 
of limited training data availability while maintaining high-quality, task-relevant outputs.

**Keywords:** Language Models, Transfer Learning, Synthetic Data, Fine-Tuning, {domain.title()}, {base_model}

---

## 1. Introduction

### 1.1 Background

Large language models (LLMs) have demonstrated remarkable capabilities across diverse natural language processing tasks. 
However, adapting these models to specific domains or tasks often requires substantial amounts of high-quality training data, 
which can be expensive, time-consuming, or difficult to obtain while maintaining privacy and ethical standards.

### 1.2 Motivation

The primary motivation for this work is to address the data scarcity problem in domain-specific language model development. 
Our specific use case—{task_description.lower()}—requires specialized knowledge and conversational patterns that may not 
be adequately represented in general-purpose language models.

### 1.3 Contributions

This work makes the following contributions:

1. **Synthetic Data Generation Framework**: We develop and apply a domain-specific synthetic data generation approach 
   that creates high-quality training examples without requiring manual annotation.

2. **Efficient Fine-Tuning**: We demonstrate effective fine-tuning of {base_model} using a relatively small dataset 
   of {dataset_size} examples, showcasing the efficiency of modern transfer learning approaches.

3. **Practical Application**: We provide a complete, production-ready model for {task_description.lower()} that can 
   be deployed immediately or serve as a foundation for further specialization.

---

## 2. Related Work

### 2.1 Transfer Learning in NLP

Transfer learning has become the dominant paradigm in natural language processing, with pre-trained models like GPT, 
BERT, and their variants achieving state-of-the-art results across numerous benchmarks. Our work builds on this 
foundation by demonstrating efficient domain adaptation.

### 2.2 Synthetic Data Generation

Recent work has shown that synthetic data can effectively augment or even replace human-annotated data for specific tasks. 
Our approach extends these findings to conversational AI and domain-specific language generation.

### 2.3 Domain Adaptation

Domain adaptation techniques allow models trained on one domain to perform well on another. Our work contributes to 
this area by combining synthetic data generation with fine-tuning for efficient domain-specific model creation.

---

## 3. Methodology

### 3.1 Base Model Selection

We selected **{base_model}** as our base model for the following reasons:

- **Architecture**: Modern transformer-based architecture with proven generation capabilities
- **Size**: Appropriate balance between capability and computational efficiency
- **Compatibility**: Well-supported by the HuggingFace ecosystem
- **Performance**: Strong baseline performance on general language tasks

### 3.2 Synthetic Data Generation

{methodology_notes if methodology_notes else f'''Our synthetic data generation process consists of several key components:

**Domain Knowledge Base:**
We curated domain-specific vocabulary, concepts, and contexts relevant to the {domain} domain. This knowledge base 
includes:
- Key topics and terminology
- Common question-answer patterns
- Domain-specific use cases
- Contextual scenarios

**Template-Based Generation:**
We employed template-based generation with intelligent variable substitution:
- Multiple conversation templates
- Dynamic topic and concept insertion
- Natural language variation
- Context-appropriate responses

**Quality Assurance:**
Each generated example undergoes validation:
- Coherence checking
- Domain relevance verification
- Diversity analysis
- Edge case inclusion'''}

### 3.3 Training Configuration

Our training setup utilized the following hyperparameters:

| Parameter | Value |
|-----------|-------|
| Base Model | {base_model} |
| Training Examples | {dataset_size} |
| Epochs | {training_params.get('epochs', 'N/A')} |
| Learning Rate | {training_params.get('learning_rate', 'N/A')} |
| Batch Size | {training_params.get('batch_size', 'N/A')} |
| Gradient Accumulation | {training_params.get('gradient_accumulation', 4)} steps |
| Optimizer | AdamW |
| Precision | Mixed (FP16) |

**Training Procedure:**

1. **Data Preparation**: Synthetic examples were tokenized using the base model's tokenizer
2. **Model Initialization**: Started from pre-trained {base_model} weights
3. **Fine-Tuning**: Applied supervised fine-tuning with causal language modeling objective
4. **Optimization**: Used gradient accumulation for memory efficiency
5. **Validation**: Monitored training loss for convergence

### 3.4 Implementation Details

Our implementation leverages:
- **Framework**: HuggingFace Transformers
- **Training Tool**: Architech AI Model Architect
- **Infrastructure**: Cloud-based GPU/CPU resources
- **Optimization**: Automatic mixed precision training

---

## 4. Results

### 4.1 Training Outcomes

{results_summary if results_summary else f'''The model successfully converged during training, demonstrating:

- **Stable Training**: Loss decreased consistently across epochs
- **No Overfitting**: Training remained stable without signs of overfitting to the small dataset
- **Efficient Learning**: Model adapted to domain-specific patterns effectively

**Qualitative Observations:**
- Generated text shows strong alignment with the {domain} domain
- Model produces coherent, contextually appropriate responses
- Task-specific vocabulary and concepts are properly utilized
- Conversation flow is natural and relevant to intended use case'''}

### 4.2 Model Capabilities

The fine-tuned model demonstrates:

1. **Domain Expertise**: Strong understanding of {domain}-specific concepts
2. **Task Alignment**: Outputs are well-aligned with {task_description.lower()}
3. **Coherence**: Generated text maintains logical consistency
4. **Flexibility**: Adapts to various prompts within the domain

### 4.3 Limitations

We acknowledge the following limitations:

- **Dataset Size**: With {dataset_size} examples, coverage of edge cases may be limited
- **Synthetic Origin**: Training data may not capture all real-world nuances
- **Domain Specificity**: Performance may degrade on out-of-domain inputs
- **Evaluation**: Comprehensive quantitative evaluation remains future work

---

## 5. Discussion

### 5.1 Effectiveness of Synthetic Data

Our results demonstrate that synthetically generated data can effectively fine-tune language models for specific tasks. 
The quality of outputs suggests that carefully designed synthetic data can capture essential patterns needed for 
domain adaptation.

### 5.2 Practical Implications

This work has several practical implications:

- **Accessibility**: Reduces barriers to creating custom language models
- **Privacy**: Eliminates need for potentially sensitive real-world data
- **Efficiency**: Enables rapid prototyping and iteration
- **Scalability**: Framework can be applied to diverse domains and tasks

### 5.3 Future Directions

Several promising directions for future work include:

1. **Quantitative Evaluation**: Comprehensive benchmarking against domain-specific metrics
2. **Dataset Scaling**: Investigation of performance vs. dataset size trade-offs
3. **Hybrid Approaches**: Combining synthetic and real data for enhanced performance
4. **Multi-Domain Transfer**: Exploring transfer learning across related domains

---

## 6. Conclusion

We presented **{model_name}**, a fine-tuned language model for {task_description.lower()}, demonstrating the 
effectiveness of synthetic data generation for domain-specific model adaptation. Our approach successfully created 
a specialized model using {dataset_size} synthetically generated examples, proving that efficient domain adaptation 
is achievable without large-scale manual data collection.

The model shows strong task alignment and domain expertise, validating our methodology. This work contributes to 
the growing body of evidence that synthetic data, when carefully designed, can serve as an effective alternative 
or complement to human-annotated data for language model fine-tuning.

As language models continue to evolve, techniques for efficient, ethical, and accessible model adaptation will 
become increasingly important. Our work provides a practical framework for creating custom language models that 
can be applied across diverse domains and use cases.

---

## 7. References

1. HuggingFace Transformers: State-of-the-art Natural Language Processing
2. Attention Is All You Need (Vaswani et al., 2017)
3. Language Models are Few-Shot Learners (Brown et al., 2020)
4. Transfer Learning in Natural Language Processing (Ruder, 2019)

---

## Appendix A: Model Architecture

**Base Architecture:** {base_model}

The model inherits the transformer-based architecture of the base model, with all parameters fine-tuned for the 
specific task.

## Appendix B: Training Logs

Training completed successfully with stable convergence. Detailed logs available in model repository.

## Appendix C: Code Availability

Model and code are available at: https://huggingface.co/your-username/{model_name}

---

## Acknowledgments

This research was conducted using Architech AI Model Architect, an open-source tool for automated language model 
development. We thank the HuggingFace team for providing the infrastructure and tools that made this work possible.

---

**Contact:** For questions about this work, please open an issue in the model repository.

**Date:** {timestamp}  
**Version:** 1.0

---

*This paper was automatically generated by Architech AI Model Architect. Please review and customize as needed for publication.*
"""
        
        # Save paper
        paper_path = self.templates_dir / f"{model_name}_research_paper.md"
        with open(paper_path, 'w') as f:
            f.write(paper)
        
        return paper, str(paper_path)
    
    def generate_both_documents(
        self,
        model_name: str,
        task_description: str,
        base_model: str,
        dataset_size: int,
        num_epochs: int,
        learning_rate: float,
        batch_size: int,
        domain: str = "general",
        intended_use: str = "",
        limitations: str = "",
        methodology_notes: str = "",
        results_summary: str = "",
        progress=gr.Progress()
    ) -> Tuple[str, str, str, str]:
        """Generate both model card and research paper"""
        
        progress(0.3, "📝 Generating Model Card...")
        
        training_params = {
            'epochs': num_epochs,
            'learning_rate': learning_rate,
            'batch_size': batch_size,
            'gradient_accumulation': 4,
            'hardware': 'GPU/CPU (auto-detected)'
        }
        
        model_card, card_path = self.generate_model_card(
            model_name, task_description, base_model, dataset_size,
            training_params, domain, intended_use, limitations
        )
        
        progress(0.7, "📄 Generating Research Paper...")
        
        paper, paper_path = self.generate_research_paper(
            model_name, task_description, base_model, dataset_size,
            training_params, domain, methodology_notes, results_summary
        )
        
        progress(1.0, "✅ Documentation Generated!")
        
        return model_card, card_path, paper, paper_path

doc_generator = DocumentationGenerator()# ==================== GRADIO INTERFACE ====================

def create_gradio_interface():
    agent = ArchitechAgent()
    
    with gr.Blocks(title="🏗️ Architech", theme=gr.themes.Soft()) as demo:
        gr.Markdown("""
        # 🏗️ **Architech - Your AI Model Architect**
        
        *Describe what you want, and I'll build it for you!*
        """)
        
        with gr.Tabs():
            # Generate Dataset
            with gr.Tab("📊 Generate Dataset"):
                with gr.Row():
                    with gr.Column():
                        task_desc = gr.Textbox(label="Task Description", lines=3, 
                            placeholder="E.g., 'Customer support chatbot for tech products'")
                        domain = gr.Dropdown(
                            choices=["technology", "healthcare", "finance", "general"],
                            label="Domain", value="general")
                        dataset_size = gr.Slider(50, 500, 100, step=50, label="Dataset Size")
                        format_type = gr.Dropdown(
                            choices=["conversational", "instruction"],
                            label="Format", value="conversational")
                        gen_btn = gr.Button("🎨 Generate Dataset", variant="primary")
                    with gr.Column():
                        gen_output = gr.Markdown()
                
                gen_btn.click(
                    fn=agent.generate_synthetic_dataset_wrapper,
                    inputs=[task_desc, domain, dataset_size, format_type, gr.State("medium")],
                    outputs=gen_output
                )
            
            # Train Model
            with gr.Tab("🚀 Train Model"):
                with gr.Row():
                    with gr.Column():
                        task_desc_train = gr.Textbox(label="Task Description", lines=2)
                        model_name = gr.Textbox(label="Model Name", placeholder="my-awesome-model")
                        hf_token = gr.Textbox(label="HuggingFace Token", type="password")
                        use_synthetic = gr.Checkbox(label="Generate New Synthetic Data", value=True)
                        
                        with gr.Group(visible=False) as dataset_group:
                            gr.Markdown("### 📊 Select Existing Dataset")
                            dataset_dropdown = gr.Dropdown(
                                label="Choose Dataset",
                                choices=[],
                                interactive=True
                            )
                            refresh_datasets_btn = gr.Button("🔄 Refresh Datasets", size="sm")
                            dataset_preview = gr.Markdown()
                            
                            def refresh_dataset_list():
                                datasets = dataset_manager.list_available_datasets()
                                choices = [name for name, path in datasets]
                                return gr.Dropdown(choices=choices)
                            
                            def show_dataset_preview(dataset_name):
                                if dataset_name:
                                    datasets = dataset_manager.list_available_datasets()
                                    for name, path in datasets:
                                        if name == dataset_name:
                                            return dataset_manager.get_dataset_preview(path)
                                return "Select a dataset to preview"
                            
                            refresh_datasets_btn.click(
                                fn=refresh_dataset_list,
                                outputs=dataset_dropdown
                            )
                            
                            dataset_dropdown.change(
                                fn=show_dataset_preview,
                                inputs=dataset_dropdown,
                                outputs=dataset_preview
                            )
                        
                        with gr.Group(visible=False) as custom_data_group:
                            training_data_input = gr.Textbox(
                                label="Training Data (one example per line) OR Dataset Path",
                                placeholder="Human: Hello\nAssistant: Hi!\n\nOR: ./synthetic_datasets/synthetic_general_conversational_20260126.json",
                                lines=8
                            )
                        
                        # Toggle visibility
                        def toggle_data_source(use_synth):
                            return gr.update(visible=not use_synth), gr.update(visible=not use_synth)
                        
                        use_synthetic.change(
                            fn=toggle_data_source,
                            inputs=use_synthetic,
                            outputs=[dataset_group, custom_data_group]
                        )
                        
                        with gr.Accordion("⚙️ Advanced", open=False):
                            base_model = gr.Dropdown(
                                choices=["distilgpt2", "gpt2", "microsoft/DialoGPT-small"],
                                label="Base Model", value="distilgpt2")
                            learning_rate = gr.Slider(1e-5, 5e-4, 2e-4, label="Learning Rate")
                            num_epochs = gr.Slider(1, 5, 3, step=1, label="Epochs")
                            batch_size = gr.Slider(1, 4, 2, step=1, label="Batch Size")
                        
                        train_btn = gr.Button("🎯 Train Model", variant="primary")
                    
                    with gr.Column():
                        train_output = gr.Markdown()
                
                def prepare_training_data(use_synth, dataset_name, custom_data):
                    """Prepare training data based on selection"""
                    if use_synth:
                        return ""  # Will generate new data
                    elif dataset_name:
                        # Use selected dataset
                        datasets = dataset_manager.list_available_datasets()
                        for name, path in datasets:
                            if name == dataset_name:
                                return path
                    return custom_data
                
                train_btn.click(
                    fn=lambda task, dataset_name, custom, model, token, base, synth, lr, epochs, batch: agent.train_custom_model(
                        task,
                        prepare_training_data(synth, dataset_name, custom),
                        model,
                        token,
                        base,
                        synth,
                        gr.State("general"),
                        gr.State(100),
                        lr,
                        epochs,
                        batch
                    ),
                    inputs=[
                        task_desc_train, dataset_dropdown, training_data_input,
                        model_name, hf_token, base_model, use_synthetic,
                        learning_rate, num_epochs, batch_size
                    ],
                    outputs=train_output
                )
            
            # Test Model
            with gr.Tab("🧪 Test Model"):
                with gr.Row():
                    with gr.Column():
                        test_model_name = gr.Textbox(label="Model Name", 
                            placeholder="username/model-name")
                        test_token = gr.Textbox(label="HuggingFace Token", type="password")
                        load_btn = gr.Button("📥 Load Model")
                        
                        gr.Markdown("---")
                        
                        test_prompt = gr.Textbox(label="Test Prompt", lines=3,
                            placeholder="Enter your prompt here...")
                        max_length = gr.Slider(50, 200, 100, label="Max Length")
                        temperature = gr.Slider(0.1, 1.0, 0.7, label="Temperature")
                        
                        test_btn = gr.Button("🎯 Generate", variant="primary")
                    
                    with gr.Column():
                        load_output = gr.Markdown()
                        test_output = gr.Markdown()
                
                load_btn.click(
                    fn=model_inference.load_model,
                    inputs=[test_model_name, test_token],
                    outputs=load_output
                )
                
                test_btn.click(
                    fn=model_inference.generate_text,
                    inputs=[test_model_name, test_prompt, max_length, temperature, gr.State(0.9)],
                    outputs=test_output
                )
            
            # Documentation Generation Tab
            with gr.Tab("📄 Generate Documentation"):
                gr.Markdown("""
                ### Generate Professional Model Card & Research Paper
                Automatically create comprehensive documentation for your models
                """)
                
                with gr.Row():
                    with gr.Column():
                        gr.Markdown("### 📋 Model Information")
                        
                        doc_model_name = gr.Textbox(
                            label="Model Name",
                            placeholder="my-awesome-model"
                        )
                        
                        doc_task_desc = gr.Textbox(
                            label="Task Description",
                            placeholder="Customer support chatbot for technical products",
                            lines=2
                        )
                        
                        doc_base_model = gr.Dropdown(
                            choices=["distilgpt2", "gpt2", "microsoft/DialoGPT-small", "other"],
                            label="Base Model",
                            value="distilgpt2"
                        )
                        
                        with gr.Row():
                            doc_dataset_size = gr.Number(
                                label="Dataset Size",
                                value=100,
                                precision=0
                            )
                            doc_domain = gr.Dropdown(
                                choices=["technology", "healthcare", "finance", "education", "general"],
                                label="Domain",
                                value="general"
                            )
                        
                        with gr.Row():
                            doc_epochs = gr.Number(label="Epochs", value=3, precision=0)
                            doc_lr = gr.Number(label="Learning Rate", value=0.0002)
                            doc_batch = gr.Number(label="Batch Size", value=2, precision=0)
                        
                        with gr.Accordion("📝 Optional Details", open=False):
                            doc_intended_use = gr.Textbox(
                                label="Intended Use (optional)",
                                placeholder="Describe specific use cases...",
                                lines=3
                            )
                            
                            doc_limitations = gr.Textbox(
                                label="Known Limitations (optional)",
                                placeholder="Describe any known limitations...",
                                lines=3
                            )
                            
                            doc_methodology = gr.Textbox(
                                label="Methodology Notes (optional)",
                                placeholder="Additional methodology details...",
                                lines=3
                            )
                            
                            doc_results = gr.Textbox(
                                label="Results Summary (optional)",
                                placeholder="Summary of model performance...",
                                lines=3
                            )
                        
                        generate_docs_btn = gr.Button("📄 Generate Documentation", variant="primary", size="lg")
                    
                    with gr.Column():
                        gr.Markdown("### 📥 Generated Documents")
                        
                        doc_status = gr.Markdown("*Generate documents to see preview*")
                        
                        with gr.Tabs():
                            with gr.Tab("📋 Model Card"):
                                model_card_output = gr.Markdown()
                                model_card_file = gr.File(label="Download Model Card")
                            
                            with gr.Tab("📄 Research Paper"):
                                paper_output = gr.Markdown()
                                paper_file = gr.File(label="Download Research Paper")
                
                def generate_and_display_docs(
                    name, task, base, size, domain, epochs, lr, batch,
                    intended, limitations, methodology, results, progress=gr.Progress()
                ):
                    try:
                        model_card, card_path, paper, paper_path = doc_generator.generate_both_documents(
                            name, task, base, int(size), int(epochs), float(lr), int(batch),
                            domain, intended, limitations, methodology, results, progress
                        )
                        
                        status = f"""✅ **Documentation Generated Successfully!**

📋 **Model Card:** `{Path(card_path).name}`
📄 **Research Paper:** `{Path(paper_path).name}`

**Files saved to:** `./generated_docs/`

**What's Next?**
1. Review the documents in the tabs above
2. Download and customize if needed
3. Upload to your model repository on HuggingFace
4. Share with the community!
"""
                        
                        # Truncate for preview
                        card_preview = model_card[:5000] + "\n\n*... (truncated for preview, download for full content)*" if len(model_card) > 5000 else model_card
                        paper_preview = paper[:5000] + "\n\n*... (truncated for preview, download for full content)*" if len(paper) > 5000 else paper
                        
                        return status, card_preview, card_path, paper_preview, paper_path
                    
                    except Exception as e:
                        error_msg = f"❌ Error generating documentation: {str(e)}"
                        return error_msg, "", None, "", None
                
                generate_docs_btn.click(
                    fn=generate_and_display_docs,
                    inputs=[
                        doc_model_name, doc_task_desc, doc_base_model,
                        doc_dataset_size, doc_domain, doc_epochs, doc_lr, doc_batch,
                        doc_intended_use, doc_limitations, doc_methodology, doc_results
                    ],
                    outputs=[doc_status, model_card_output, model_card_file, paper_output, paper_file]
                )
                
                gr.Markdown("""
                ---
                ### 💡 Documentation Tips
                
                **Model Card:**
                - Standard format recognized by HuggingFace
                - Includes model details, training info, and usage examples
                - Ready to upload to your model repository
                
                **Research Paper:**
                - Academic-style documentation
                - Describes methodology and approach
                - Great for sharing your work formally
                
                **Best Practices:**
                - Fill in optional fields for more detailed documentation
                - Customize generated docs before publishing
                - Keep documentation up-to-date with model changes
                - Include ethical considerations and limitations
                """)
            
            # Repository Chat Tab
            with gr.Tab("💬 Repository Chat"):
                gr.Markdown("""
                ### Chat with Your HuggingFace Repositories
                Manage your models and datasets conversationally!
                """)
                
                with gr.Row():
                    with gr.Column():
                        repo_token = gr.Textbox(
                            label="HuggingFace Token",
                            type="password",
                            placeholder="hf_..."
                        )
                        init_btn = gr.Button("🔐 Initialize Session", variant="primary")
                        init_output = gr.Markdown()
                        
                        init_btn.click(
                            fn=lambda token: repo_chat.initialize_session(token)[1],
                            inputs=repo_token,
                            outputs=init_output
                        )
                
                gr.Markdown("---")
                
                with gr.Row():
                    with gr.Column(scale=2):
                        chatbot = gr.Chatbot(
                            label="Repository Assistant",
                            height=400
                        )
                        
                        with gr.Row():
                            chat_input = gr.Textbox(
                                label="Message",
                                placeholder="Try: 'List my models' or 'Show my datasets'",
                                scale=4
                            )
                            send_btn = gr.Button("Send", variant="primary", scale=1)
                        
                        gr.Markdown("""
                        **Quick Commands:**
                        - "List my models" - Show all your models
                        - "Show my datasets" - Show all your datasets  
                        - "Info about [model-name]" - Get model details
                        - "Help" - See all commands
                        """)
                    
                    with gr.Column(scale=1):
                        gr.Markdown("### 🗑️ Delete Repository")
                        
                        delete_repo_id = gr.Textbox(
                            label="Repository ID",
                            placeholder="username/model-name"
                        )
                        
                        delete_repo_type = gr.Radio(
                            choices=["model", "dataset"],
                            label="Type",
                            value="model"
                        )
                        
                        delete_repo_btn = gr.Button("🗑️ Delete", variant="stop")
                        delete_repo_output = gr.Markdown()
                        
                        delete_repo_btn.click(
                            fn=repo_chat.delete_repo,
                            inputs=[delete_repo_id, delete_repo_type],
                            outputs=delete_repo_output
                        )
                
                def chat_respond(message, history):
                    if not message.strip():
                        return history, ""
                    
                    bot_response = repo_chat.chat_with_repos(message)
                    history.append((message, bot_response))
                    return history, ""