import os from typing import List, Dict from datetime import datetime import torch from torch import nn import gradio as gr import pandas as pd from datasets import Dataset from transformers import ( AutoModelForCausalLM, AutoTokenizer, GenerationConfig, ) from peft import LoraConfig, get_peft_model from trl import DPOConfig, DPOTrainer # ========================================================= # MODEL LIST (from your BIAS demo) # ========================================================= MODEL_CHOICES = [ # Very small / light (good for CPU Spaces) "distilgpt2", "gpt2", "sshleifer/tiny-gpt2", "LiquidAI/LFM2-350M", "google/gemma-3-270m-it", "Qwen/Qwen2.5-0.5B-Instruct", "mkurman/NeuroBLAST-V3-SYNTH-EC-150000", # Small–medium (~1–2B) – still reasonable on CPU, just slower "TinyLlama/TinyLlama-1.1B-Chat-v1.0", "google/gemma-3-1b-it", "meta-llama/Llama-3.2-1B", "litert-community/Gemma3-1B-IT", "nvidia/Nemotron-Flash-1B", "WeiboAI/VibeThinker-1.5B", "Qwen/Qwen3-1.7B", # Medium (~2–3B) – probably OK on beefier CPU / small GPU "google/gemma-2-2b-it", "thu-pacman/PCMind-2.1-Kaiyuan-2B", "opendatalab/MinerU-HTML", "ministral/Ministral-3b-instruct", "HuggingFaceTB/SmolLM3-3B", "meta-llama/Llama-3.2-3B-Instruct", "nvidia/Nemotron-Flash-3B-Instruct", "Qwen/Qwen2.5-3B-Instruct", # Heavier (4–8B) – you really want a GPU Space for these "Qwen/Qwen3-4B", "Qwen/Qwen3-4B-Thinking-2507", "Qwen/Qwen3-4B-Instruct-2507", "mistralai/Mistral-7B-Instruct-v0.2", "allenai/Olmo-3-7B-Instruct", "Qwen/Qwen2.5-7B-Instruct", "meta-llama/Meta-Llama-3-8B-Instruct", "meta-llama/Llama-3.1-8B", "meta-llama/Llama-3.1-8B-Instruct", "openbmb/MiniCPM4.1-8B", "deepseek-ai/DeepSeek-R1-Distill-Llama-8B", "rl-research/DR-Tulu-8B", ] DEFAULT_MODEL = "Qwen/Qwen2.5-0.5B-Instruct" TRAINED_MODEL_DIR = "trained_model" # ========================================================= # GLOBALS & CONFIG # ========================================================= device = "cuda" if torch.cuda.is_available() else "cpu" tokenizer = None policy_model = None ref_model = None DEFAULT_DPO_CONFIG = DPOConfig( beta=0.1, output_dir="dpo_demo", num_train_epochs=1, per_device_train_batch_size=1, per_device_eval_batch_size=1, remove_unused_columns=False, logging_steps=1, gradient_accumulation_steps=1, learning_rate=1e-4, evaluation_strategy="no", # warning is fine with current versions warmup_steps=0, fp16=False, save_steps=0, report_to="none", ) # ========================================================= # LORA TARGET-MODULE HELPER # ========================================================= def guess_lora_target_modules(model_name: str, base_model) -> List[str]: """ Heuristically choose good LoRA target modules based on the model type/name. - GPT-2-like: use c_attn/c_proj - LLaMA/Gemma/Mistral/Qwen/etc: use q/k/v/o + MLP projections - Fallback: scan Linear module names for known patterns """ model_type = getattr(base_model.config, "model_type", "") or "" name_lower = model_name.lower() # GPT-2 / DistilGPT-2 / Tiny GPT-2 if ( "gpt2" in model_type or "gpt2" in name_lower or "tiny-gpt2" in name_lower or "distilgpt2" in name_lower ): return ["c_attn", "c_proj"] # LLaMA / Gemma / Mistral / Qwen / Olmo / MiniCPM / SmolLM / Nemotron etc. if any( t in model_type for t in [ "llama", "gemma", "mistral", "qwen", "qwen2", "olmo", "minicpm", "smollm", "nemotron", ] ): return ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] # Fallback: inspect Linear modules and see what’s there linear_leaf_names = [] for name, module in base_model.named_modules(): if isinstance(module, nn.Linear): linear_leaf_names.append(name.split(".")[-1]) candidates = [ "q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", "c_attn", "c_proj", ] found = sorted(set(n for n in candidates if n in linear_leaf_names)) if found: return found # If absolutely nothing matches, bail with a clear error raise ValueError( f"Could not guess LoRA target modules for model '{model_name}' " f"(model_type='{model_type}'). " f"Try setting target_modules manually for this model." ) # ========================================================= # MODEL LOADING # ========================================================= def load_base_model(model_name: str) -> str: """ Load tokenizer + base model, then create: - policy_model: LoRA-adapted (trainable) - ref_model: frozen base model for DPO """ global tokenizer, policy_model, ref_model tokenizer = AutoTokenizer.from_pretrained( model_name, trust_remote_code=True, ) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token tokenizer.padding_side = "right" base_model = AutoModelForCausalLM.from_pretrained( model_name, trust_remote_code=True, ) base_model.config.use_cache = False base_model.config.pad_token_id = tokenizer.eos_token_id # Choose LoRA target modules dynamically target_modules = guess_lora_target_modules(model_name, base_model) peft_config = LoraConfig( r=4, target_modules=target_modules, task_type="CAUSAL_LM", lora_alpha=8, lora_dropout=0.1, bias="none", ) # Policy model = base + LoRA (trainable) policy = get_peft_model(base_model, peft_config) policy.to(device) policy.eval() # Reference model = frozen base model reference = AutoModelForCausalLM.from_pretrained( model_name, trust_remote_code=True, ) reference.config.use_cache = False reference.config.pad_token_id = tokenizer.eos_token_id reference.to(device) for p in reference.parameters(): p.requires_grad = False reference.eval() policy_model = policy ref_model = reference return ( f"Loaded base model: **{model_name}** on **{device}** " f"with LoRA target_modules={target_modules}" ) # Load default on startup initial_status = load_base_model(DEFAULT_MODEL) # ========================================================= # UTILS # ========================================================= def build_generation_config( do_sample: bool, temperature: float, max_new_tokens: int, top_k: int = 20, top_p: float = 0.9, ) -> GenerationConfig: """ Helper to build a GenerationConfig from UI settings. """ # Clamp values a bit just to be safe temperature = max(0.0, float(temperature)) max_new_tokens = int(max_new_tokens) return GenerationConfig( do_sample=bool(do_sample), temperature=temperature, top_k=top_k, top_p=top_p, max_new_tokens=max_new_tokens, pad_token_id=tokenizer.eos_token_id, ) def generate_text( model: nn.Module, prompt: str, gen_config: GenerationConfig, style_prefix: str = "", ) -> str: model.eval() full_prompt = style_prefix + prompt inputs = tokenizer( full_prompt, return_tensors="pt", padding=False, ).to(device) with torch.no_grad(): outputs = model.generate( **inputs, do_sample=gen_config.do_sample, top_k=gen_config.top_k, top_p=gen_config.top_p, temperature=gen_config.temperature, max_new_tokens=gen_config.max_new_tokens, pad_token_id=gen_config.pad_token_id, ) text = tokenizer.decode(outputs[0], skip_special_tokens=True) if text.startswith(full_prompt): return text[len(full_prompt):].strip() return text.strip() def preferences_to_df(preferences: List[Dict]) -> pd.DataFrame: if not preferences: return pd.DataFrame(columns=["prompt", "chosen", "rejected"]) return pd.DataFrame(preferences) def list_trained_model_files() -> List[str]: """ Return a list of filepaths under TRAINED_MODEL_DIR (for download). """ if not os.path.isdir(TRAINED_MODEL_DIR): return [] files: List[str] = [] for root, dirs, filenames in os.walk(TRAINED_MODEL_DIR): for name in filenames: files.append(os.path.join(root, name)) return files # ========================================================= # DPO CALLBACKS # ========================================================= def generate_candidates( prompt: str, do_sample: bool, temperature: float, max_new_tokens: int, ) -> tuple[str, str]: """ Generate Answer A (balanced) and Answer B (creative-ish), using the same core generation settings from the GUI. """ if not prompt.strip(): return "", "" # Build two configs from the same UI settings, # but make B slightly more "wild" by bumping top_k / temperature a bit balanced_config = build_generation_config( do_sample=do_sample, temperature=temperature, max_new_tokens=max_new_tokens, top_k=20, top_p=0.9, ) # For creative answer, nudge temperature and top_k a bit, but still # keep them tied to UI settings. creative_temp = float(temperature) + 0.4 creative_config = build_generation_config( do_sample=do_sample, temperature=creative_temp, max_new_tokens=max_new_tokens, top_k=50, top_p=0.95, ) style_balanced = ( "You are a helpful, careful assistant. " "Answer clearly and sensibly.\n\nUser: " ) style_creative = ( "You are a creative assistant who explores unusual ideas and stronger opinions, " "while still staying safe.\n\nUser: " ) answer_a = generate_text( policy_model, prompt, balanced_config, style_prefix=style_balanced, ) answer_b = generate_text( policy_model, prompt, creative_config, style_prefix=style_creative, ) return answer_a, answer_b def save_preference( prompt: str, answer_a: str, answer_b: str, custom_answer: str, preference_mode: str, state_preferences: List[Dict], ): """ Encode a preference in one of four ways: - Prefer A over B -> chosen=A, rejected=B - Prefer B over A -> chosen=B, rejected=A - Prefer custom over A -> chosen=custom, rejected=A - Prefer custom over B -> chosen=custom, rejected=B """ msg = "" if not prompt.strip(): msg = "No prompt provided." return state_preferences, preferences_to_df(state_preferences), msg if not answer_a.strip() or not answer_b.strip(): msg = "Generate both model answers before saving a preference." return state_preferences, preferences_to_df(state_preferences), msg if not preference_mode: msg = "Please choose how to encode the preference." return state_preferences, preferences_to_df(state_preferences), msg preference_mode = preference_mode.strip() chosen = None rejected = None if preference_mode == "Prefer A over B": chosen = answer_a rejected = answer_b elif preference_mode == "Prefer B over A": chosen = answer_b rejected = answer_a elif preference_mode == "Prefer custom over A": if not custom_answer.strip(): msg = "You selected 'Prefer custom over A' but did not provide a custom answer." return state_preferences, preferences_to_df(state_preferences), msg chosen = custom_answer rejected = answer_a elif preference_mode == "Prefer custom over B": if not custom_answer.strip(): msg = "You selected 'Prefer custom over B' but did not provide a custom answer." return state_preferences, preferences_to_df(state_preferences), msg chosen = custom_answer rejected = answer_b else: msg = f"Unknown preference mode: {preference_mode}" return state_preferences, preferences_to_df(state_preferences), msg entry = { "prompt": prompt.strip(), "chosen": chosen.strip(), "rejected": rejected.strip(), } state_preferences = list(state_preferences) + [entry] df = preferences_to_df(state_preferences) msg = f"Saved preference #{len(state_preferences)}." return state_preferences, df, msg def train_dpo_model( state_preferences: List[Dict], num_epochs: int, learning_rate: float, beta: float, progress=gr.Progress(track_tqdm=True), ): """ Run DPO training on the accumulated preferences. Shows a progress bar/spinner and returns: - a detailed status message - a 'last trained' timestamp string - a list of saved model files for download """ global policy_model, ref_model progress(0.0, desc="Checking preferences...") if not state_preferences: return ( "⚠️ No preferences collected yet. Add some first.", "**Last trained:** never", [], ) dataset = Dataset.from_list(state_preferences) progress(0.2, desc="Configuring DPO trainer...") dpo_config = DPOConfig( **{ **DEFAULT_DPO_CONFIG.to_dict(), "num_train_epochs": int(num_epochs), "learning_rate": float(learning_rate), "beta": float(beta), } ) trainer = DPOTrainer( model=policy_model, ref_model=ref_model, args=dpo_config, train_dataset=dataset, eval_dataset=None, tokenizer=tokenizer, max_length=256, ) progress(0.4, desc="Training model with DPO...") trainer.train() progress(0.75, desc="Finalizing and moving model to device...") policy_model = trainer.model policy_model.to(device) policy_model.eval() # Save the trained model + tokenizer so you can download them progress(0.9, desc="Saving trained model to disk...") os.makedirs(TRAINED_MODEL_DIR, exist_ok=True) policy_model.save_pretrained(TRAINED_MODEL_DIR) tokenizer.save_pretrained(TRAINED_MODEL_DIR) files = list_trained_model_files() progress(1.0, desc="Done") n = len(state_preferences) finished_at = datetime.now().strftime("%Y-%m-%d %H:%M:%S") msg = f"""### ✅ Training complete - Preference pairs used: **{n}** - Epochs: **{num_epochs}** - Learning rate: **{learning_rate}** - DPO beta (strength): **{beta}** The tuned policy model + tokenizer have been saved to `{TRAINED_MODEL_DIR}/`. You can download them using the file list below. """ last_trained_msg = f"**Last trained:** {finished_at}" return msg, last_trained_msg, files def generate_from_aligned_model( prompt: str, do_sample: bool, temperature: float, max_new_tokens: int, ) -> str: if not prompt.strip(): return "" gen_config = build_generation_config( do_sample=do_sample, temperature=temperature, max_new_tokens=max_new_tokens, top_k=20, top_p=0.9, ) style_balanced = ( "You are a helpful, careful assistant. " "Answer clearly and sensibly.\n\nUser: " ) return generate_text( policy_model, prompt, gen_config, style_prefix=style_balanced, ) def on_model_change( model_name: str, _state_preferences: List[Dict], ): """ When the user picks a new base model: - reload tokenizer + policy_model + ref_model - clear collected preferences (since they belong to previous model) - reset training status, 'last trained', and download list """ status = load_base_model(model_name) empty_prefs: List[Dict] = [] df = preferences_to_df(empty_prefs) reset_msg = ( status + "\n\nPreferences cleared (new model = new preference data)." ) last_trained_reset = "**Last trained:** (reset for new base model)" files_reset: List[str] = [] # returns: model_status, prefs, pref_table_df, train_status, last_trained, files return reset_msg, empty_prefs, df, "", last_trained_reset, files_reset # ========================================================= # GRADIO UI # ========================================================= with gr.Blocks() as demo: gr.Markdown( """ # 🔧 DPO Playground – Preference Tuning on Different Models - Pick a **base model** from the dropdown. - Ask a question and generate two answers: - **A** = balanced / normal - **B** = creative / more extreme - Optionally write **your own ideal answer**. - Choose how to encode the preference (e.g. A over B, custom over A, etc.). - Collect several preferences and **train the model with DPO**. - Test how the aligned policy model behaves on new prompts. - Download the tuned model (LoRA adapter + tokenizer) after training. - **Control temperature, sampling, and max_new_tokens directly in the UI.** """ ) state_preferences = gr.State([]) with gr.Row(): model_dropdown = gr.Dropdown( choices=MODEL_CHOICES, value=DEFAULT_MODEL, label="Base model", ) model_status = gr.Markdown(initial_status) # ----------------------------------------------------- # Collect preferences tab # ----------------------------------------------------- with gr.Tab("Collect preferences"): with gr.Row(): prompt_input = gr.Textbox( label="Prompt", placeholder="Ask anything...", lines=3, ) gr.Markdown("### Generation settings for Answer A & B") with gr.Row(): gen_do_sample = gr.Checkbox( value=True, label="Use sampling (do_sample)", ) gen_temperature = gr.Slider( minimum=0.0, maximum=1.5, value=0.8, step=0.05, label="Temperature", ) gen_max_new_tokens = gr.Slider( minimum=4, maximum=256, value=128, step=4, label="Max new tokens", ) generate_btn = gr.Button("Generate A & B") with gr.Row(): answer_a_box = gr.Textbox( label="Answer A (balanced / normal)", lines=8, ) answer_b_box = gr.Textbox( label="Answer B (creative / more extreme)", lines=8, ) custom_answer_box = gr.Textbox( label="Your own ideal answer (optional)", lines=8, placeholder="If you want, write the answer you *wish* the model had given.", ) preference_mode = gr.Radio( choices=[ "Prefer A over B", "Prefer B over A", "Prefer custom over A", "Prefer custom over B", ], label="How should this preference be encoded?", ) save_pref_btn = gr.Button("Save preference") pref_status = gr.Markdown("") pref_table = gr.Dataframe( headers=["prompt", "chosen", "rejected"], label="Collected preferences (for DPO training)", wrap=True, ) generate_btn.click( fn=generate_candidates, inputs=[prompt_input, gen_do_sample, gen_temperature, gen_max_new_tokens], outputs=[answer_a_box, answer_b_box], ) save_pref_btn.click( fn=save_preference, inputs=[ prompt_input, answer_a_box, answer_b_box, custom_answer_box, preference_mode, state_preferences, ], outputs=[ state_preferences, pref_table, pref_status, ], ) # ----------------------------------------------------- # Train & test tab # ----------------------------------------------------- with gr.Tab("Train & test DPO model"): gr.Markdown( "Train the LoRA-adapted policy model using your preferences " "with **Direct Preference Optimization (DPO)**." ) with gr.Row(): num_epochs_slider = gr.Slider( minimum=1, maximum=5, step=1, value=1, label="Number of epochs", ) lr_slider = gr.Slider( minimum=1e-5, maximum=5e-4, step=1e-5, value=1e-4, label="Learning rate", ) beta_slider = gr.Slider( minimum=0.05, maximum=0.5, step=0.05, value=0.1, label="DPO beta (strength)", ) train_btn = gr.Button("Train DPO model", variant="primary") train_status = gr.Markdown("") last_trained = gr.Markdown("**Last trained:** never") download_files = gr.Files( label="Trained model files (adapter + tokenizer)", interactive=False, ) train_btn.click( fn=train_dpo_model, inputs=[ state_preferences, num_epochs_slider, lr_slider, beta_slider, ], outputs=[train_status, last_trained, download_files], ) gr.Markdown("## Try the current policy model") with gr.Row(): test_do_sample = gr.Checkbox( value=False, label="Use sampling (do_sample) for test", ) test_temperature = gr.Slider( minimum=0.0, maximum=1.5, value=0.0, step=0.05, label="Temperature (test)", ) test_max_new_tokens = gr.Slider( minimum=4, maximum=256, value=64, step=4, label="Max new tokens (test)", ) test_prompt = gr.Textbox( label="Test prompt", placeholder="Ask something to see the aligned model...", lines=3, ) test_btn = gr.Button("Generate from DPO policy model") test_answer = gr.Textbox( label="Policy model answer", lines=8, ) test_btn.click( fn=generate_from_aligned_model, inputs=[ test_prompt, test_do_sample, test_temperature, test_max_new_tokens, ], outputs=test_answer, ) # model change: reload + clear prefs + reset train status + last trained + downloads model_dropdown.change( fn=on_model_change, inputs=[model_dropdown, state_preferences], outputs=[ model_status, state_preferences, pref_table, train_status, last_trained, download_files, ], ) if __name__ == "__main__": demo.queue().launch()