Spaces:
Running
Running
| """ | |
| Alkaid A — Hugging Face Spaces App | |
| Training + Inference UI hosted entirely on Hugging Face | |
| Space Type: Docker (GPU required — A10G or A100 recommended) | |
| """ | |
| import os | |
| import json | |
| import threading | |
| import time | |
| import gradio as gr | |
| from pathlib import Path | |
| # ============================================================================= | |
| # CONFIGURATION | |
| # ============================================================================= | |
| BASE_MODEL = "Jackrong/Qwen3.5-27B-Claude-4.6-Opus-Reasoning-Distilled" | |
| OUTPUT_DIR = "/home/user/app/alkaid_a_checkpoints" | |
| FINAL_DIR = "/home/user/app/alkaid_a_final" | |
| DATA_DIR = "/home/user/app/data" | |
| LOG_FILE = "/home/user/app/training.log" | |
| SYSTEM_PROMPT = ( | |
| "You are Alkaid A, an advanced AI coding and deployment assistant. " | |
| "You follow a rigorous multi-phase workflow: (1) Provide detailed feedback " | |
| "with pros/cons on code or plans, identifying weak points and breaks. " | |
| "(2) Guide through a detailed debug phase. (3) Outline a deployment strategy " | |
| "ready for production. (4) Repeat debugging with variations across five iterations. " | |
| "(5) Conduct deep dives on integration covering security, scalability, and compliance. " | |
| "(6) Test all API endpoints and set monitoring. (7) Scrape help docs, check tool " | |
| "compatibility, infer issues, and adjust on the fly. (8) Ensure every version is " | |
| "backed up in a GitHub releases folder, starting at 00.00.00 and incrementing by " | |
| "00.00.01. (9) Guide through pushing changes to a GitHub repository. " | |
| "(10) Include user testing, performance benchmarking, and hardening. " | |
| "(11) Add documentation for future developers and automated testing. " | |
| "(12) Summarize what went well and acknowledge progress." | |
| ) | |
| # Global state | |
| training_status = {"running": False, "progress": "", "log": ""} | |
| loaded_model = {"model": None, "tokenizer": None} | |
| # ============================================================================= | |
| # TRAINING TAB | |
| # ============================================================================= | |
| def format_opus_example(example): | |
| """Convert Opus dataset row to chat format.""" | |
| assistant_content = f"<think>\n{example['thinking']}\n</think>\n\n{example['solution']}" | |
| return { | |
| "messages": [ | |
| {"role": "system", "content": SYSTEM_PROMPT}, | |
| {"role": "user", "content": example["problem"]}, | |
| {"role": "assistant", "content": assistant_content}, | |
| ] | |
| } | |
| def run_training( | |
| hf_token, | |
| hub_repo_id, | |
| learning_rate, | |
| num_epochs, | |
| lora_rank, | |
| max_seq_length, | |
| batch_size, | |
| use_4bit, | |
| custom_data_text, | |
| progress=gr.Progress(track_tqdm=True), | |
| ): | |
| """Run the full training pipeline.""" | |
| global training_status | |
| if training_status["running"]: | |
| return "Training is already running. Please wait." | |
| training_status["running"] = True | |
| training_status["log"] = "" | |
| def log(msg): | |
| training_status["log"] += msg + "\n" | |
| training_status["progress"] = msg | |
| print(msg) | |
| try: | |
| # --- Login --- | |
| log("Step 1/7: Authenticating with Hugging Face...") | |
| if hf_token: | |
| from huggingface_hub import login | |
| login(token=hf_token) | |
| log(" Logged in successfully.") | |
| else: | |
| log(" No token provided — will save locally only.") | |
| # --- Load Model --- | |
| log(f"Step 2/7: Loading base model ({BASE_MODEL})...") | |
| from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig | |
| import torch | |
| if use_4bit: | |
| bnb_config = BitsAndBytesConfig( | |
| load_in_4bit=True, | |
| bnb_4bit_quant_type="nf4", | |
| bnb_4bit_compute_dtype=torch.bfloat16, | |
| bnb_4bit_use_double_quant=True, | |
| ) | |
| model = AutoModelForCausalLM.from_pretrained( | |
| BASE_MODEL, | |
| quantization_config=bnb_config, | |
| device_map="auto", | |
| torch_dtype=torch.bfloat16, | |
| trust_remote_code=True, | |
| ) | |
| else: | |
| model = AutoModelForCausalLM.from_pretrained( | |
| BASE_MODEL, | |
| device_map="auto", | |
| torch_dtype=torch.bfloat16, | |
| trust_remote_code=True, | |
| ) | |
| tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True) | |
| if tokenizer.pad_token is None: | |
| tokenizer.pad_token = tokenizer.eos_token | |
| log(" Model loaded.") | |
| # --- Attach LoRA --- | |
| log(f"Step 3/7: Attaching LoRA (rank={lora_rank})...") | |
| from peft import LoraConfig, get_peft_model, TaskType | |
| lora_config = LoraConfig( | |
| r=int(lora_rank), | |
| lora_alpha=int(lora_rank), | |
| target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], | |
| lora_dropout=0.0, | |
| bias="none", | |
| task_type=TaskType.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()) | |
| log(f" LoRA attached. Trainable: {trainable:,} / {total:,} ({100*trainable/total:.2f}%)") | |
| # --- Prepare Data --- | |
| log("Step 4/7: Preparing training data...") | |
| from datasets import load_dataset, Dataset, concatenate_datasets | |
| # Load Opus dataset | |
| opus_ds = load_dataset("nohurry/Opus-4.6-Reasoning-3000x-filtered", split="train") | |
| opus_ds = opus_ds.map(format_opus_example, remove_columns=opus_ds.column_names) | |
| log(f" Opus dataset: {len(opus_ds)} examples") | |
| # Process custom data if provided | |
| if custom_data_text and custom_data_text.strip(): | |
| custom_rows = [] | |
| for line in custom_data_text.strip().split("\n"): | |
| line = line.strip() | |
| if line: | |
| try: | |
| custom_rows.append(json.loads(line)) | |
| except json.JSONDecodeError: | |
| log(f" Warning: Skipping invalid JSON line") | |
| if custom_rows: | |
| custom_ds = Dataset.from_list(custom_rows) | |
| # Weight custom data 3x | |
| all_ds = concatenate_datasets([opus_ds, custom_ds, custom_ds, custom_ds]) | |
| log(f" Custom data: {len(custom_rows)} examples (weighted 3x)") | |
| else: | |
| all_ds = opus_ds | |
| else: | |
| all_ds = opus_ds | |
| # Apply chat template | |
| def apply_template(example): | |
| text = tokenizer.apply_chat_template( | |
| example["messages"], tokenize=False, add_generation_prompt=False | |
| ) | |
| return {"text": text} | |
| all_ds = all_ds.map(apply_template) | |
| all_ds = all_ds.shuffle(seed=42) | |
| log(f" Total training examples: {len(all_ds)}") | |
| # --- Train --- | |
| log(f"Step 5/7: Training ({num_epochs} epochs, lr={learning_rate})...") | |
| from trl import SFTTrainer, SFTConfig | |
| training_args = SFTConfig( | |
| output_dir=OUTPUT_DIR, | |
| per_device_train_batch_size=int(batch_size), | |
| gradient_accumulation_steps=4, | |
| warmup_steps=10, | |
| num_train_epochs=int(num_epochs), | |
| learning_rate=float(learning_rate), | |
| optim="adamw_8bit", | |
| lr_scheduler_type="cosine", | |
| bf16=True, | |
| fp16=False, | |
| logging_steps=5, | |
| save_steps=50, | |
| save_total_limit=2, | |
| max_seq_length=int(max_seq_length), | |
| dataset_text_field="text", | |
| report_to="none", | |
| seed=42, | |
| ) | |
| trainer = SFTTrainer( | |
| model=model, | |
| tokenizer=tokenizer, | |
| train_dataset=all_ds, | |
| args=training_args, | |
| ) | |
| result = trainer.train() | |
| log(f" Training complete! Final loss: {result.training_loss:.4f}") | |
| # --- Save --- | |
| log("Step 6/7: Saving model...") | |
| model.save_pretrained(FINAL_DIR) | |
| tokenizer.save_pretrained(FINAL_DIR) | |
| log(f" Saved to {FINAL_DIR}") | |
| # --- Push to Hub --- | |
| if hf_token and hub_repo_id: | |
| log(f"Step 7/7: Pushing to {hub_repo_id}...") | |
| model.push_to_hub(hub_repo_id, use_auth_token=hf_token) | |
| tokenizer.push_to_hub(hub_repo_id, use_auth_token=hf_token) | |
| log(f" Live at: https://huggingface.co/{hub_repo_id}") | |
| else: | |
| log("Step 7/7: Skipped push (no token or repo ID).") | |
| log("\n" + "=" * 50) | |
| log("TRAINING COMPLETE!") | |
| log("=" * 50) | |
| training_status["running"] = False | |
| return training_status["log"] | |
| except Exception as e: | |
| training_status["running"] = False | |
| error_msg = f"\nERROR: {str(e)}" | |
| log(error_msg) | |
| import traceback | |
| log(traceback.format_exc()) | |
| return training_status["log"] | |
| # ============================================================================= | |
| # INFERENCE TAB | |
| # ============================================================================= | |
| def load_model_for_inference(model_source, hf_token): | |
| """Load model for inference (either local checkpoint or from Hub).""" | |
| global loaded_model | |
| try: | |
| from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig | |
| import torch | |
| source = FINAL_DIR if model_source == "Local checkpoint" else model_source | |
| if hf_token: | |
| from huggingface_hub import login | |
| login(token=hf_token) | |
| bnb_config = BitsAndBytesConfig( | |
| load_in_4bit=True, | |
| bnb_4bit_quant_type="nf4", | |
| bnb_4bit_compute_dtype=torch.bfloat16, | |
| ) | |
| loaded_model["tokenizer"] = AutoTokenizer.from_pretrained(source, trust_remote_code=True) | |
| loaded_model["model"] = AutoModelForCausalLM.from_pretrained( | |
| source, | |
| quantization_config=bnb_config, | |
| device_map="auto", | |
| torch_dtype=torch.bfloat16, | |
| trust_remote_code=True, | |
| ) | |
| if loaded_model["tokenizer"].pad_token is None: | |
| loaded_model["tokenizer"].pad_token = loaded_model["tokenizer"].eos_token | |
| return f"Model loaded from: {source}" | |
| except Exception as e: | |
| return f"Error loading model: {str(e)}" | |
| def generate_response(user_message, temperature, max_tokens, system_override): | |
| """Generate a response using the loaded model.""" | |
| if loaded_model["model"] is None: | |
| return "Please load a model first using the 'Load Model' button." | |
| try: | |
| import torch | |
| model = loaded_model["model"] | |
| tokenizer = loaded_model["tokenizer"] | |
| sys_prompt = system_override if system_override.strip() else SYSTEM_PROMPT | |
| messages = [ | |
| {"role": "system", "content": sys_prompt}, | |
| {"role": "user", "content": user_message}, | |
| ] | |
| input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt").to(model.device) | |
| with torch.no_grad(): | |
| output = model.generate( | |
| input_ids, | |
| max_new_tokens=int(max_tokens), | |
| temperature=float(temperature), | |
| top_p=0.9, | |
| do_sample=True, | |
| pad_token_id=tokenizer.pad_token_id, | |
| ) | |
| response = tokenizer.decode(output[0][input_ids.shape[-1]:], skip_special_tokens=True) | |
| return response | |
| except Exception as e: | |
| return f"Generation error: {str(e)}" | |
| # ============================================================================= | |
| # GRADIO UI | |
| # ============================================================================= | |
| # Load default custom data | |
| default_custom_data = "" | |
| custom_data_path = Path("/home/user/app/alkaid_a_training_data.jsonl") | |
| if custom_data_path.exists(): | |
| default_custom_data = custom_data_path.read_text() | |
| with gr.Blocks( | |
| title="Alkaid A — Train & Deploy", | |
| theme=gr.themes.Soft(primary_hue="indigo", neutral_hue="slate"), | |
| css=""" | |
| .main-title { text-align: center; margin-bottom: 0; } | |
| .subtitle { text-align: center; color: #6b7280; margin-top: 4px; } | |
| """ | |
| ) as app: | |
| gr.Markdown("# Alkaid A", elem_classes="main-title") | |
| gr.Markdown( | |
| "*Fine-tuned from Qwen3.5-27B-Claude-4.6-Opus-Reasoning-Distilled*", | |
| elem_classes="subtitle" | |
| ) | |
| with gr.Tabs(): | |
| # ---- TRAINING TAB ---- | |
| with gr.Tab("Train", id="train"): | |
| gr.Markdown("### Train Your Model") | |
| gr.Markdown( | |
| "This will fine-tune the base model using LoRA on the Opus reasoning " | |
| "dataset plus your custom data, then push the result to your Hugging Face repo." | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| hf_token_train = gr.Textbox( | |
| label="Hugging Face Token (Write access)", | |
| type="password", | |
| placeholder="hf_...", | |
| ) | |
| hub_repo = gr.Textbox( | |
| label="Hub Repo ID (e.g., YourName/Alkaid-A)", | |
| placeholder="YourUsername/Alkaid-A", | |
| ) | |
| with gr.Column(scale=1): | |
| with gr.Row(): | |
| lr = gr.Number(label="Learning Rate", value=2e-4) | |
| epochs = gr.Number(label="Epochs", value=3, precision=0) | |
| with gr.Row(): | |
| rank = gr.Number(label="LoRA Rank", value=16, precision=0) | |
| seq_len = gr.Number(label="Max Seq Length", value=2048, precision=0) | |
| with gr.Row(): | |
| bs = gr.Number(label="Batch Size", value=1, precision=0) | |
| fourbit = gr.Checkbox(label="4-bit Quantization", value=True) | |
| custom_data = gr.Code( | |
| label="Custom Training Data (JSONL — one JSON object per line)", | |
| value=default_custom_data, | |
| language="json", | |
| lines=10, | |
| ) | |
| train_btn = gr.Button("Start Training", variant="primary", size="lg") | |
| train_output = gr.Textbox(label="Training Log", lines=20, interactive=False) | |
| train_btn.click( | |
| fn=run_training, | |
| inputs=[hf_token_train, hub_repo, lr, epochs, rank, seq_len, bs, fourbit, custom_data], | |
| outputs=train_output, | |
| ) | |
| # ---- INFERENCE TAB ---- | |
| with gr.Tab("Chat", id="chat"): | |
| gr.Markdown("### Chat with Alkaid A") | |
| with gr.Row(): | |
| model_source = gr.Textbox( | |
| label="Model Source", | |
| value="Local checkpoint", | |
| placeholder="Local checkpoint OR HuggingFace repo ID", | |
| ) | |
| hf_token_infer = gr.Textbox( | |
| label="HF Token (if loading from Hub)", | |
| type="password", | |
| placeholder="hf_...", | |
| ) | |
| load_btn = gr.Button("Load Model") | |
| load_status = gr.Textbox(label="Status", interactive=False) | |
| load_btn.click(fn=load_model_for_inference, inputs=[model_source, hf_token_infer], outputs=load_status) | |
| system_box = gr.Textbox( | |
| label="System Prompt (optional override)", | |
| value="", | |
| placeholder="Leave empty to use default Alkaid A system prompt", | |
| lines=3, | |
| ) | |
| chatbot_input = gr.Textbox( | |
| label="Your Message", | |
| placeholder="Paste your code or describe your plan...", | |
| lines=6, | |
| ) | |
| with gr.Row(): | |
| temp = gr.Slider(label="Temperature", minimum=0.1, maximum=1.5, value=0.7, step=0.1) | |
| max_tok = gr.Slider(label="Max Tokens", minimum=256, maximum=4096, value=2048, step=256) | |
| gen_btn = gr.Button("Generate", variant="primary") | |
| response_box = gr.Textbox(label="Alkaid A Response", lines=20, interactive=False) | |
| gen_btn.click( | |
| fn=generate_response, | |
| inputs=[chatbot_input, temp, max_tok, system_box], | |
| outputs=response_box, | |
| ) | |
| # ---- ABOUT TAB ---- | |
| with gr.Tab("About", id="about"): | |
| gr.Markdown(""" | |
| ### Alkaid A | |
| **Base Model:** [Jackrong/Qwen3.5-27B-Claude-4.6-Opus-Reasoning-Distilled](https://huggingface.co/Jackrong/Qwen3.5-27B-Claude-4.6-Opus-Reasoning-Distilled) | |
| **Training Dataset:** [nohurry/Opus-4.6-Reasoning-3000x-filtered](https://huggingface.co/datasets/nohurry/Opus-4.6-Reasoning-3000x-filtered) (2,326 reasoning examples) | |
| **Method:** LoRA SFT with 4-bit quantization | |
| **Alkaid A's Workflow:** | |
| 1. Detailed code/plan feedback with pros and cons | |
| 2. Guided debug phase | |
| 3. Production deployment strategy | |
| 4. 5x debug iterations with variations | |
| 5. Security, scalability, compliance deep dive | |
| 6. API endpoint testing and monitoring setup | |
| 7. Help doc scraping and compatibility checks | |
| 8. GitHub versioned releases (00.00.XX) | |
| 9. Guided repository push | |
| 10. User testing, benchmarking, hardening | |
| 11. Developer documentation and automated tests | |
| 12. Progress summary and acknowledgment | |
| **License:** Apache 2.0 | |
| """) | |
| if __name__ == "__main__": | |
| app.launch(server_name="0.0.0.0", server_port=7860, share=False) | |