Upload app.py with huggingface_hub

app.py

@@ -1,549 +1,21 @@
 import gradio as gr
 import torch
-import json
-import time
 import os
 
-# Ensure Gradio cache directory exists
 os.makedirs("/tmp/gradio", exist_ok=True)
-os.environ["GRADIO_TEMP_DIR"] = "/tmp/gradio"
 
-# Dataset configs - verified PUBLIC datasets
-DATASETS = {
-    "NousResearch Hermes": {
-        "name": "NousResearch/hermes-function-calling-v1",
-        "type": "tool_calling"
-    },
-    "Glaive Code Assistant": {
-        "name": "glaiveai/glaive-code-assistant-v2",
-        "type": "tool_calling"
-    },
-    "OpenAssistant": {
-        "name": "OpenAssistant/oasst1",
-        "type": "chat"
-    },
-    "Custom Upload": {
-        "name": "custom",
-        "type": "custom"
-    }
-}
+def test_gpu():
+    if torch.cuda.is_available():
+        return f"GPU: {torch.cuda.get_device_name(0)}\nVRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB"
+    return "No GPU found"
 
-# Model configs - organized by size
-MODELS = {
-    # Tiny (< 1B) - Fast training
-    "Qwen2.5-0.5B-Instruct": "Qwen/Qwen2.5-0.5B-Instruct",
-    "SmolLM2-360M": "HuggingFaceTB/SmolLM2-360M-Instruct",
+with gr.Blocks(title="VertexElite AI") as demo:
+    gr.Markdown("# VertexElite AI Tool Trainer")
+    gr.Markdown("Testing T4 GPU...")
 
-    # Small (1-2B) - Good balance
-    "Qwen2.5-1.5B-Instruct": "Qwen/Qwen2.5-1.5B-Instruct",
-    "TinyLlama-1.1B": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
-    "Gemma-2-2B": "google/gemma-2-2b-it",
-    "SmolLM2-1.7B": "HuggingFaceTB/SmolLM2-1.7B-Instruct",
-    "Llama-3.2-1B": "meta-llama/Llama-3.2-1B-Instruct",
+    btn = gr.Button("Check GPU")
+    output = gr.Textbox(label="GPU Info")
 
-    # Medium (3-4B) - More capable
-    "Qwen2.5-3B-Instruct": "Qwen/Qwen2.5-3B-Instruct",
-    "Llama-3.2-3B": "meta-llama/Llama-3.2-3B-Instruct",
-    "Phi-3.5-mini": "microsoft/Phi-3.5-mini-instruct",
+    btn.click(test_gpu, outputs=output)
 
-    # Large (7B+) - Most capable (slower)
-    "Qwen2.5-7B-Instruct": "Qwen/Qwen2.5-7B-Instruct",
-    "Mistral-7B-v0.3": "mistralai/Mistral-7B-Instruct-v0.3",
-    "Gemma-2-9B": "google/gemma-2-9b-it",
-}
-
-def format_example(example, dataset_type):
-    """Format dataset examples for training"""
-    if dataset_type == "tool_calling":
-        if "conversations" in example:
-            convs = example["conversations"]
-            if isinstance(convs, list):
-                text = ""
-                for conv in convs:
-                    role = conv.get("from", conv.get("role", "user"))
-                    content = conv.get("value", conv.get("content", ""))
-                    if role in ["system", "gpt", "assistant"]:
-                        role_tag = "assistant" if role in ["gpt", "assistant"] else "system"
-                    else:
-                        role_tag = "user"
-                    text += f"<|im_start|>{role_tag}\n{content}\n<|im_end|>\n"
-                return text
-
-        query = example.get("question", example.get("instruction", example.get("query", "")))
-        response = example.get("response", example.get("answer", example.get("output", "")))
-        return f"<|im_start|>user\n{query}\n<|im_end|>\n<|im_start|>assistant\n{response}\n<|im_end|>"
-
-    elif dataset_type == "chat":
-        if "messages" in example:
-            text = ""
-            for msg in example["messages"]:
-                role = msg.get("role", "user")
-                content = msg.get("content", msg.get("text", ""))
-                text += f"<|im_start|>{role}\n{content}\n<|im_end|>\n"
-            return text
-        return str(example)
-
-    elif dataset_type == "custom":
-        if "text" in example:
-            return example["text"]
-        elif "messages" in example:
-            text = ""
-            for msg in example["messages"]:
-                role = msg.get("role", "user")
-                content = msg.get("content", "")
-                text += f"<|im_start|>{role}\n{content}\n<|im_end|>\n"
-            return text
-        elif "instruction" in example:
-            instr = example.get("instruction", "")
-            inp = example.get("input", "")
-            out = example.get("output", "")
-            prompt = f"{instr}\n{inp}".strip()
-            return f"<|im_start|>user\n{prompt}\n<|im_end|>\n<|im_start|>assistant\n{out}\n<|im_end|>"
-        return str(example)
-
-    return str(example)
-
-def load_custom_dataset(file_path):
-    """Load custom dataset from uploaded file - supports json, jsonl, csv, parquet, txt"""
-    data = []
-    ext = file_path.lower().split('.')[-1] if '.' in file_path else ''
-
-    try:
-        # Parquet
-        if ext == 'parquet':
-            import pandas as pd
-            df = pd.read_parquet(file_path)
-            data = df.to_dict('records')
-
-        # CSV
-        elif ext == 'csv':
-            import pandas as pd
-            df = pd.read_csv(file_path)
-            data = df.to_dict('records')
-
-        # Text file (one example per line)
-        elif ext == 'txt':
-            with open(file_path, 'r', encoding='utf-8') as f:
-                for line in f:
-                    line = line.strip()
-                    if line:
-                        data.append({"text": line})
-
-        # JSON/JSONL
-        else:
-            with open(file_path, 'r', encoding='utf-8') as f:
-                content = f.read().strip()
-                if content.startswith('{'):
-                    # JSONL format
-                    for line in content.split('\n'):
-                        line = line.strip()
-                        if line:
-                            try:
-                                data.append(json.loads(line))
-                            except:
-                                pass
-                elif content.startswith('['):
-                    # JSON array
-                    data = json.loads(content)
-
-    except Exception as e:
-        print(f"Error loading custom dataset: {e}")
-
-    return data
-
-def train_model(model_name, datasets, max_samples, lora_r, lora_alpha, epochs, output_name, custom_file, continue_training, continue_from, progress=gr.Progress()):
-    """Train model with LoRA"""
-
-    logs = []
-    def log(msg):
-        logs.append(f"[{time.strftime('%H:%M:%S')}] {msg}")
-        return "\n".join(logs)
-
-    try:
-        yield log("Checking training libraries...")
-
-        from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, BitsAndBytesConfig, Trainer, DataCollatorForLanguageModeling
-        from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, PeftModel
-        from datasets import load_dataset, Dataset
-
-        hf_token = os.environ.get("HF_TOKEN")
-
-        # Determine base model
-        if continue_training and continue_from:
-            yield log(f"Continuing training from: {continue_from}")
-            # Load adapter config to find base model
-            try:
-                from huggingface_hub import hf_hub_download
-                import json as json_lib
-
-                config_path = hf_hub_download(
-                    repo_id=continue_from,
-                    filename="adapter_config.json",
-                    token=hf_token
-                )
-                with open(config_path) as f:
-                    adapter_config = json_lib.load(f)
-                model_id = adapter_config.get("base_model_name_or_path", MODELS.get(model_name, model_name))
-                yield log(f"Base model: {model_id}")
-            except Exception as e:
-                yield log(f"Could not load adapter config: {e}")
-                model_id = MODELS.get(model_name, model_name)
-        else:
-            model_id = MODELS.get(model_name, model_name)
-
-        yield log(f"Loading model: {model_id}")
-
-        bnb_config = BitsAndBytesConfig(
-            load_in_4bit=True,
-            bnb_4bit_quant_type="nf4",
-            bnb_4bit_compute_dtype=torch.bfloat16,
-            bnb_4bit_use_double_quant=True
-        )
-
-        tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-        if tokenizer.pad_token is None:
-            tokenizer.pad_token = tokenizer.eos_token
-
-        model = AutoModelForCausalLM.from_pretrained(
-            model_id,
-            quantization_config=bnb_config,
-            device_map="auto",
-            trust_remote_code=True
-        )
-
-        yield log(f"Model loaded on {next(model.parameters()).device}")
-
-        # Apply LoRA or load existing adapter
-        if continue_training and continue_from:
-            yield log(f"Loading existing LoRA adapter from {continue_from}...")
-            try:
-                model = PeftModel.from_pretrained(
-                    model,
-                    continue_from,
-                    is_trainable=True,
-                    token=hf_token
-                )
-                yield log("Existing adapter loaded - continuing training!")
-            except Exception as e:
-                yield log(f"Could not load adapter: {e}")
-                yield log("Starting fresh with new LoRA...")
-                model = prepare_model_for_kbit_training(model)
-                lora_config = LoraConfig(
-                    r=int(lora_r),
-                    lora_alpha=int(lora_alpha),
-                    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
-                    lora_dropout=0.05,
-                    bias="none",
-                    task_type="CAUSAL_LM"
-                )
-                model = get_peft_model(model, lora_config)
-        else:
-            yield log(f"Applying new LoRA (r={lora_r}, alpha={lora_alpha})")
-            model = prepare_model_for_kbit_training(model)
-            lora_config = LoraConfig(
-                r=int(lora_r),
-                lora_alpha=int(lora_alpha),
-                target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
-                lora_dropout=0.05,
-                bias="none",
-                task_type="CAUSAL_LM"
-            )
-            model = get_peft_model(model, lora_config)
-
-        trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
-        total = sum(p.numel() for p in model.parameters())
-        yield log(f"Trainable: {trainable:,} / {total:,} ({100*trainable/total:.2f}%)")
-
-        yield log(f"Loading datasets: {datasets}")
-
-        all_texts = []
-        samples_per_dataset = max_samples // max(len(datasets), 1)
-
-        for ds_name in datasets:
-            if ds_name == "Custom Upload":
-                if custom_file is not None:
-                    yield log("Loading custom dataset...")
-                    try:
-                        custom_data = load_custom_dataset(custom_file)
-                        if custom_data:
-                            for example in custom_data[:samples_per_dataset]:
-                                text = format_example(example, "custom")
-                                all_texts.append(text)
-                            yield log(f"Loaded {min(len(custom_data), samples_per_dataset)} samples from custom dataset")
-                    except Exception as e:
-                        yield log(f"Error: {e}")
-                continue
-
-            if ds_name not in DATASETS:
-                yield log(f"Unknown dataset: {ds_name}")
-                continue
-
-            ds_config = DATASETS[ds_name]
-            yield log(f"Loading {ds_name}...")
-
-            try:
-                ds = load_dataset(ds_config["name"], split="train", trust_remote_code=True)
-                sample_size = min(samples_per_dataset, len(ds))
-                ds = ds.shuffle(seed=42).select(range(sample_size))
-
-                for example in ds:
-                    text = format_example(example, ds_config["type"])
-                    all_texts.append(text)
-
-                yield log(f"Loaded {sample_size} samples from {ds_name}")
-
-            except Exception as e:
-                yield log(f"Error loading {ds_name}: {e}")
-
-        if not all_texts:
-            yield log("ERROR: No data loaded!")
-            return
-
-        yield log(f"Total training samples: {len(all_texts)}")
-
-        # Tokenize
-        yield log("Tokenizing dataset...")
-
-        tokenized = tokenizer(all_texts, truncation=True, max_length=512, padding=True)
-        train_dataset = Dataset.from_dict({
-            "input_ids": tokenized["input_ids"],
-            "attention_mask": tokenized["attention_mask"],
-            "labels": tokenized["input_ids"].copy()
-        })
-
-        yield log("Setting up trainer...")
-
-        training_args = TrainingArguments(
-            output_dir=f"./outputs/{output_name}",
-            num_train_epochs=int(epochs),
-            per_device_train_batch_size=1,
-            gradient_accumulation_steps=4,
-            learning_rate=2e-4,
-            warmup_ratio=0.03,
-            logging_steps=10,
-            save_strategy="epoch",
-            fp16=True,
-            optim="paged_adamw_8bit",
-            report_to="none",
-            remove_unused_columns=False
-        )
-
-        data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
-
-        trainer = Trainer(
-            model=model,
-            args=training_args,
-            train_dataset=train_dataset,
-            data_collator=data_collator
-        )
-
-        yield log("Starting training...")
-
-        trainer.train()
-
-        yield log("Training complete!")
-        yield log(f"Saving model to ./outputs/{output_name}")
-
-        trainer.save_model(f"./outputs/{output_name}")
-        tokenizer.save_pretrained(f"./outputs/{output_name}")
-
-        # Push to HuggingFace Hub
-        yield log("Pushing to HuggingFace Hub...")
-        try:
-            from huggingface_hub import HfApi, create_repo, login
-
-            if hf_token:
-                login(token=hf_token)
-
-            repo_id = f"crittiksglobal/{output_name}"
-
-            try:
-                create_repo(repo_id, repo_type="model", exist_ok=True, token=hf_token)
-            except:
-                pass
-
-            api = HfApi(token=hf_token)
-            api.upload_folder(
-                folder_path=f"./outputs/{output_name}",
-                repo_id=repo_id,
-                repo_type="model"
-            )
-
-            yield log(f"Model pushed to: https://huggingface.co/{repo_id}")
-            yield log("You can download from there!")
-
-        except Exception as e:
-            yield log(f"Push failed: {e}")
-
-        yield log("\n=== Training Complete ===")
-        yield log(f"Base model: {model_id}")
-        yield log(f"Samples: {len(all_texts)}")
-        yield log(f"LoRA rank: {lora_r}")
-        if continue_training:
-            yield log(f"Continued from: {continue_from}")
-        yield log(f"Download: https://huggingface.co/crittiksglobal/{output_name}")
-
-    except Exception as e:
-        import traceback
-        yield log(f"ERROR: {e}\n{traceback.format_exc()}")
-
-def test_model(prompt, output_name):
-    """Test a trained model"""
-    try:
-        from transformers import AutoModelForCausalLM, AutoTokenizer
-        from peft import PeftModel
-
-        hf_token = os.environ.get("HF_TOKEN")
-
-        # Try local first, then Hub
-        model_path = f"./outputs/{output_name}"
-        if not os.path.exists(model_path):
-            model_path = f"crittiksglobal/{output_name}"
-
-        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True, token=hf_token)
-        model = AutoModelForCausalLM.from_pretrained(
-            model_path,
-            device_map="auto",
-            torch_dtype=torch.float16,
-            trust_remote_code=True,
-            token=hf_token
-        )
-
-        formatted = f"<|im_start|>user\n{prompt}\n<|im_end|>\n<|im_start|>assistant\n"
-        inputs = tokenizer(formatted, return_tensors="pt").to(model.device)
-
-        with torch.no_grad():
-            outputs = model.generate(
-                **inputs,
-                max_new_tokens=256,
-                temperature=0.7,
-                do_sample=True,
-                pad_token_id=tokenizer.pad_token_id
-            )
-
-        response = tokenizer.decode(outputs[0], skip_special_tokens=False)
-
-        if "<|im_start|>assistant" in response:
-            response = response.split("<|im_start|>assistant")[-1]
-            if "<|im_end|>" in response:
-                response = response.split("<|im_end|>")[0]
-
-        return response.strip()
-
-    except Exception as e:
-        return f"Error: {e}"
-
-# Gradio UI
-with gr.Blocks(title="VertexElite AI Tool Trainer", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("""
-    # VertexElite AI Tool Trainer
-
-    Fine-tune **tool-calling** models on ZeroGPU (A10G) with LoRA.
-    """)
-
-    with gr.Tabs():
-        with gr.TabItem("Train"):
-            with gr.Row():
-                with gr.Column():
-                    model = gr.Dropdown(
-                        choices=list(MODELS.keys()),
-                        value="Qwen2.5-0.5B-Instruct",
-                        label="Base Model"
-                    )
-
-                    gr.Markdown("### Continue Training")
-                    continue_training = gr.Checkbox(
-                        label="Continue from existing model",
-                        value=False
-                    )
-                    continue_from = gr.Textbox(
-                        label="Model to continue from",
-                        value="crittiksglobal/vertexelite-v1",
-                        placeholder="crittiksglobal/vertexelite-v1"
-                    )
-
-                    gr.Markdown("### Datasets")
-                    dataset = gr.CheckboxGroup(
-                        choices=list(DATASETS.keys()),
-                        value=["Custom Upload"],
-                        label="Select Datasets"
-                    )
-
-                    custom_file = gr.File(
-                        label="Upload dataset (jsonl, json, csv, parquet, txt)",
-                        file_types=[".jsonl", ".json", ".csv", ".parquet", ".txt"],
-                        type="filepath"
-                    )
-
-                    samples = gr.Slider(100, 5000, 1000, step=100, label="Max Samples")
-
-                    with gr.Row():
-                        lora_r = gr.Slider(4, 64, 16, step=4, label="LoRA Rank")
-                        lora_alpha = gr.Slider(8, 128, 32, step=8, label="LoRA Alpha")
-
-                    epochs = gr.Slider(1, 3, 1, step=1, label="Epochs")
-                    output = gr.Textbox(value="vertexelite-v1", label="Output Name")
-
-                    train_btn = gr.Button("Start Training", variant="primary", size="lg")
-
-                with gr.Column():
-                    logs = gr.Textbox(label="Training Logs", lines=25)
-
-            train_btn.click(
-                train_model,
-                [model, dataset, samples, lora_r, lora_alpha, epochs, output, custom_file, continue_training, continue_from],
-                logs
-            )
-
-        with gr.TabItem("Test"):
-            gr.Markdown("Test your trained model")
-            with gr.Row():
-                with gr.Column():
-                    test_model_name = gr.Textbox(value="vertexelite-v1", label="Model Name")
-                    prompt = gr.Textbox(label="Prompt", value="Hello, how can you help me?", lines=3)
-                    test_btn = gr.Button("Generate", variant="primary")
-                with gr.Column():
-                    result = gr.Textbox(label="Output", lines=10)
-
-            test_btn.click(test_model, [prompt, test_model_name], result)
-
-        with gr.TabItem("Info"):
-            gr.Markdown("""
-            ### Continue Training
-
-            Check **"Continue from existing model"** to improve an existing model:
-            1. Enter the model repo (e.g., `crittiksglobal/vertexelite-v1`)
-            2. Upload new training data
-            3. Train - it will learn from new data while keeping old knowledge!
-
-            ### Datasets
-            | Dataset | Description |
-            |---------|-------------|
-            | NousResearch Hermes | Function calling |
-            | Glaive Code Assistant | Code generation |
-            | OpenAssistant | General chat |
-            | Custom Upload | Your JSONL/JSON |
-
-            ### Supported Formats
-            | Format | Description |
-            |--------|-------------|
-            | `.jsonl` | One JSON object per line |
-            | `.json` | Array of objects |
-            | `.csv` | Columns: instruction, output (or text) |
-            | `.parquet` | Same columns as CSV |
-            | `.txt` | One example per line |
-
-            ### Data Format
-            ```json
-            {"instruction": "...", "output": "..."}
-            {"messages": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}
-            {"text": "full conversation text"}
-            ```
-
-            ### Tips
-            - Continue training = model gets smarter over time
-            - Use same output name to update the model
-            - More data + more epochs = better model
-            """)
-
-demo.launch(ssr_mode=False)
+demo.launch()