From 670809f55596f95b1dff6ba320717c52618bb263 Mon Sep 17 00:00:00 2001 From: 51616 Date: Tue, 13 May 2025 10:49:31 +0000 Subject: [PATCH] webui working (not yet two loras) --- README.md | 5 + chat_templates/google/gemma-2-2b-it.jinja | 2 +- intx_sft.py | 11 +- webui/app.py | 257 +++++++++++++++++++++- webui/templates/visualize.html | 243 +++++++++++++++++++- 5 files changed, 493 insertions(+), 25 deletions(-) diff --git a/README.md b/README.md index c4a3d99..47cbbb3 100644 --- a/README.md +++ b/README.md @@ -42,6 +42,11 @@ run python intx_sft.py configs/math_and_code.yaml --model_name_or_path=google/ge run python intx_sft.py configs/gsm8k.yaml --model_name_or_path=google/gemma-2-2b-it --num_train_epochs=5 --per_device_train_batch_size=16 --gradient_accumulation_steps=1 --per_device_eval_batch_size=8 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --num_blocks=1 --num_self_attends_per_block=16 --self_attention_widening_factor=1 --eval_steps=5000 --save_steps=5000 --learning_rate=2e-5 --neftune_noise_alpha=5 --use_light_weight_lora=True --light_weight_latent_size=512 --load_best_model_at_end=True --metric_for_best_model=pwc_loss --add_negative_prompt=False --add_repeat_prompt=False --ctx_encoder_model_name_or_path=meta-llama/Llama-3.2-11B-Vision-Instruct ``` +### Continue from a checkpoint +```bash +run python intx_sft.py configs/...yaml ... --from_pretrained_checkpoint=train_outputs/runs/May09_16-25-35_slurm0-a3nodeset-4_59459_ea85a571/checkpoint-10000/pytorch_model.bin --resume_from_checkpoint=train_outputs/runs/May09_16-25-35_slurm0-a3nodeset-4_59459_ea85a571/checkpoint-10000 +``` + ### Evaluation LongBench ```bash diff --git a/chat_templates/google/gemma-2-2b-it.jinja b/chat_templates/google/gemma-2-2b-it.jinja index e5eaf57..b9ba380 100644 --- a/chat_templates/google/gemma-2-2b-it.jinja +++ b/chat_templates/google/gemma-2-2b-it.jinja @@ -5,7 +5,7 @@ {%- else %} {%- set loop_messages = messages %} {%- endif %} -{% for message in messages %} +{% for message in loop_messages %} {% if (message['role'] == 'assistant') %} {% set role = 'model' %} {% else %} diff --git a/intx_sft.py b/intx_sft.py index ede385f..04c323b 100755 --- a/intx_sft.py +++ b/intx_sft.py @@ -253,8 +253,6 @@ def main(): ] ) - accelerator = Accelerator(deepspeed_plugins=DeepSpeedPlugin("ds_config.json")) - device = accelerator.device set_seed(training_args.seed) checkpoint_dir = training_args.resume_from_checkpoint if checkpoint_dir and not os.path.isdir(checkpoint_dir): @@ -312,7 +310,6 @@ def main(): train=True, requires_grad=ctx_args.exp_setup == ExperimentSetup.FULL_FINETUNE, peft_config=get_lora_config(model_name, **vars(lora_args)), - device=device, ) ctx_name = ctx_encoder_args.ctx_encoder_model_name_or_path if ctx_name is not None: @@ -403,7 +400,7 @@ def main(): set_format=None if ctx_args.use_sequence_packing else "pt", # streaming=data_args.streaming, ) - tokenized_ds = {} + tokenized_ds = {"train": dict(), "validation": dict(), "test": dict()} for split, ds_names in zip( ["train", "validation", "test"], [data_args.train_ds_names, data_args.val_ds_names, data_args.test_ds_names], @@ -412,9 +409,9 @@ def main(): continue streaming = (split == "train") and data_args.streaming for ds_name in ds_names: - with accelerator.main_process_first(): - ds = _get_tokenized_dataset(ds_name, split, streaming=streaming) - tokenized_ds[split] = {os.path.basename(ds_name): ds} + # with accelerator.main_process_first(): + ds = _get_tokenized_dataset(ds_name, split, streaming=streaming) + tokenized_ds[split][os.path.basename(ds_name)] = ds train_ds = tokenized_ds["train"] logging.info(f"train_ds: {train_ds}") diff --git a/webui/app.py b/webui/app.py index 6bd295f..d372b57 100644 --- a/webui/app.py +++ b/webui/app.py @@ -6,7 +6,9 @@ from glob import glob import torch import yaml from flask import Flask, abort, jsonify, render_template, request -from transformers import pipeline +from transformers import pipeline, AutoTokenizer + +from ctx_to_lora.model_loading import get_tokenizer app = Flask(__name__) @@ -15,6 +17,7 @@ chat_generator = None chat_model_name = None chat_history = None device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +modulated_model = None def get_run_data(run_path): @@ -116,6 +119,63 @@ def get_generated_text_data(run_path): return generated_data +def get_available_checkpoints(run): + """ + Finds all checkpoint directories in a run folder. + + Args: + run: The name of the training run. + + Returns: + A list of checkpoint directories. + """ + logdir = os.path.join(TRAIN_OUTPUTS_DIR, run) + if not os.path.isdir(logdir): + return [] + + checkpoints = sorted( + [d for d in glob(os.path.join(logdir, "checkpoint-*")) if os.path.isdir(d)], + key=lambda d: int(d.split("-")[-1]), + ) + + return checkpoints + + +def load_custom_chat_template(tokenizer, model_name): + """ + Loads a custom chat template for models that have issues with the built-in templates. + + Args: + tokenizer: The HuggingFace tokenizer. + model_name: The model name to find an appropriate template for. + """ + # Extract model family and name for template matching + model_parts = model_name.split("/")[-1].split("-") + model_family = model_name.split("/")[0] if "/" in model_name else "" + + # Check for Gemma models + if "gemma" in model_name.lower(): + template_path = "chat_templates/google/gemma-2-2b-it.jinja" + if os.path.exists(template_path): + with open(template_path, "r") as f: + template_content = f.read() + tokenizer.chat_template = template_content + print(f"Loaded custom chat template from {template_path}") + return True + + # Check Meta models + elif "meta-llama" in model_name.lower() or "llama" in model_name.lower(): + template_path = "chat_templates/meta-llama/llama-2-7b-chat.jinja" + if os.path.exists(template_path): + with open(template_path, "r") as f: + template_content = f.read() + tokenizer.chat_template = template_content + print(f"Loaded custom chat template from {template_path}") + return True + + return False + + @app.route("/") def index(): """ @@ -152,6 +212,10 @@ def visualize(run): reverse=True, ) + # Get available checkpoints + checkpoints = get_available_checkpoints(run) + checkpoint_names = [os.path.basename(cp) for cp in checkpoints] + selected_eval_folder = request.args.get("eval_folder") if eval_folders and selected_eval_folder: @@ -180,6 +244,7 @@ def visualize(run): config_path = os.path.join(logdir, "config.yaml") try: with open(config_path) as f: + yaml.add_constructor("!", lambda loader, node: None) config = yaml.safe_load(f) model_name = config.get("model_name_or_path") except FileNotFoundError: @@ -193,6 +258,7 @@ def visualize(run): data=data, generated_data=generated_data, model_name=model_name, + checkpoints=checkpoint_names, ) @@ -204,26 +270,96 @@ def load_model(): global chat_generator global chat_model_name global chat_history + global modulated_model print("Loading model...") + # Reset chat history chat_history = [{"role": "system", "content": ""}] + # Reset modulated model if it exists + modulated_model = None + run = request.form["run"] logdir = os.path.join(TRAIN_OUTPUTS_DIR, run) - config_path = os.path.join(logdir, "config.yaml") + config_path = os.path.join(logdir, "args.yaml") try: with open(config_path) as f: + yaml.add_constructor("!", lambda loader, node: None) + yaml.add_multi_constructor( + "tag:yaml.org,2002:python/object", + lambda loader, suffix, node: None, + Loader=yaml.SafeLoader, + ) config = yaml.safe_load(f) chat_model_name = config.get("model_name_or_path") except FileNotFoundError: chat_model_name = None if chat_model_name: + tokenizer = get_tokenizer(chat_model_name, train=False) + + # Handle system roles by loading a custom chat template if needed + load_custom_chat_template(tokenizer, chat_model_name) + + print(f"Using model: {chat_model_name}") + print( + f"Tokenizer chat template available: {hasattr(tokenizer, 'chat_template')}" + ) + chat_generator = pipeline( - "text-generation", model=chat_model_name, device=device + "text-generation", model=chat_model_name, tokenizer=tokenizer, device=device ) return jsonify({"model_name": chat_model_name}) else: - return jsonify({"error": "Model name not found in config.yaml."}) + return jsonify({"error": "Model name not found in args.yaml."}) + + +@app.route("/load_checkpoint", methods=["POST"]) +def load_checkpoint(): + """ + Loads a checkpoint and creates a modulated model. + """ + global modulated_model + global chat_history + global chat_generator + + # Reset the chat history + chat_history = [{"role": "system", "content": ""}] + # Unload the regular chat model to avoid confusion + chat_generator = None + + from ctx_to_lora.modeling_utils import ModulatedPretrainedModel + + run = request.form["run"] + checkpoint = request.form["checkpoint"] + context = request.form.get("context", "") + + logdir = os.path.join(TRAIN_OUTPUTS_DIR, run) + checkpoint_path = os.path.join(logdir, checkpoint, "pytorch_model.bin") + + try: + print(f"Loading checkpoint: {checkpoint_path}") + state_dict = torch.load(checkpoint_path) + + modulated_model = ModulatedPretrainedModel.from_state_dict( + state_dict, + train=False, + use_flash_attn=False, + ) + modulated_model = modulated_model.to(device) + modulated_model.eval() + + result = {"success": True, "message": f"Loaded checkpoint {checkpoint}"} + + if context: + result["context_processed"] = True + + return jsonify(result) + except Exception as e: + print(f"Error loading checkpoint: {str(e)}") + import traceback + + traceback.print_exc() + return jsonify({"success": False, "error": str(e)}) @app.route("/chat", methods=["POST"]) @@ -232,20 +368,101 @@ def chat(): Handles chat requests and generates responses. """ global chat_history + global modulated_model print("Chat request received") message = request.form["message"] print(f"Received message: {message}") + if chat_history is None: + chat_history = [{"role": "system", "content": ""}] + chat_history.append({"role": "user", "content": message}) - if chat_generator: + if chat_generator or modulated_model: print("Generating response...") - response = chat_generator(chat_history, max_length=2**13)[0]["generated_text"][ - -1 - ]["content"] - print(f"Response: {response}") + try: + # Handle both regular chat_generator and modulated_model + if modulated_model: + print("Using modulated model for response generation") + # Process context and generate LoRA for the response + ctx_tokenizer = None + with torch.inference_mode(), torch.amp.autocast(str(device)): + from transformers import AutoTokenizer + + # Get the context and tokenize it + context = request.form.get("context", "") + ctx_encoder_model_name_or_path = ( + modulated_model.ctx_encoder_args.ctx_encoder_model_name_or_path + or modulated_model.base_model.config.name_or_path + ) + ctx_tokenizer = get_tokenizer( + ctx_encoder_model_name_or_path, train=False + ) + + # Prepare base model inputs + base_tokenizer = get_tokenizer( + modulated_model.base_model.config.name_or_path, train=False + ) + # load_custom_chat_template( + # base_tokenizer, modulated_model.base_model.config.name_or_path + # ) + + # Tokenize the context for the ctx_encoder + from ctx_to_lora.data_utils import tokenize_ctx_text + + ctx_inputs = tokenize_ctx_text( + dict(context=[context]), + ctx_tokenizer, + ) + # print(f"Context inputs: {ctx_inputs}") + ctx_ids = torch.tensor(ctx_inputs["ctx_ids"], device=device) + ctx_attn_mask = torch.tensor( + ctx_inputs["ctx_attn_mask"], device=device + ) + + print(f"chat_history: {chat_history}") + + # Tokenize the chat history + model_inputs = base_tokenizer.apply_chat_template( + chat_history, return_tensors="pt" + ).to(device) + + # Generate response with context-modulated model + outputs = modulated_model.generate( + input_ids=model_inputs, + ctx_ids=ctx_ids, + ctx_attn_mask=ctx_attn_mask, + max_new_tokens=512, + do_sample=False, + ) + + # Decode the generated response + response = base_tokenizer.decode( + outputs[0][model_inputs.shape[1] :], skip_special_tokens=True + ) + + print(f"Modulated model response: {response}") + elif chat_generator: + response = chat_generator(chat_history, max_length=2**13)[0][ + "generated_text" + ][-1]["content"] + else: + response = "No model loaded. Please load a model first." + + print(f"Response: {response}") + chat_history.append({"role": "assistant", "content": response}) + return jsonify({"response": response}) + except Exception as e: + print(f"Error generating response: {str(e)}") + import traceback + + traceback.print_exc() + return jsonify({"response": f"Error: {str(e)}"}) else: - response = "Chat model not loaded." - return jsonify({"response": response}) + return jsonify( + { + "response": "No model loaded. Please load a base model or apply a hypernetwork first." + } + ) @app.route("/update_system_msg", methods=["POST"]) @@ -322,5 +539,23 @@ def get_config(run): return jsonify({"config": "cli_args.yaml not found."}) +@app.route("/reset_chat", methods=["POST"]) +def reset_chat(): + """ + Resets the chat history to start a new conversation. + """ + global chat_history + + print("Resetting chat history") + chat_history = [{"role": "system", "content": ""}] + + # Get the system message if it was set previously + system_msg = request.form.get("system_msg", "") + if system_msg: + chat_history[0]["content"] = system_msg + + return jsonify({"success": True, "message": "Chat history reset successfully"}) + + if __name__ == "__main__": app.run(debug=True) diff --git a/webui/templates/visualize.html b/webui/templates/visualize.html index bfac570..0688595 100644 --- a/webui/templates/visualize.html +++ b/webui/templates/visualize.html @@ -230,6 +230,76 @@ border: 1px solid #ddd; border-radius: 5px; } + + .hypernetwork-container { + margin-top: 20px; + padding: 15px; + background-color: #f0f8ff; + border: 1px solid #add8e6; + border-radius: 5px; + } + + .hypernetwork-container h3 { + margin-top: 0; + margin-bottom: 15px; + color: #0066cc; + } + + .hypernetwork-container label { + display: block; + margin-bottom: 5px; + font-weight: bold; + } + + .hypernetwork-container select, + .hypernetwork-container textarea { + width: 100%; + padding: 8px; + margin-bottom: 15px; + border: 1px solid #add8e6; + border-radius: 5px; + } + + .hypernetwork-container textarea { + height: 120px; + resize: vertical; + } + + .chat-controls { + margin-top: 10px; + display: flex; + gap: 10px; + } + + .secondary-button { + background-color: #607d8b; + color: white; + padding: 8px 16px; + border: none; + border-radius: 5px; + cursor: pointer; + transition: background-color 0.3s; + } + + .secondary-button:hover { + background-color: #455a64; + } + + .instruction-text { + margin: 0 0 15px 0; + font-style: italic; + color: #555; + font-size: 0.9em; + } + + .ui-sequence { + margin-bottom: 10px; + padding: 10px; + background-color: #f5f5f5; + border-left: 4px solid #4CAF50; + color: #333; + font-size: 0.9em; + } {% endblock %} @@ -333,16 +403,52 @@ selected_eval_folder }}{% else %}root{% endif %}.

{% endif %} + +
+

How to use this interface:

+
    +
  1. First, click "Load Chat with {{ model_name }}" to initialize the base model
    2. +
  2. You can then chat directly with the base model
    3. +
  3. Alternatively, you can select a hypernetwork checkpoint to use context-informed responses
    4. +
  4. To reset the conversation at any time, click the "Reset Chat" button
    5. +
+
+
+ + + {% if checkpoints %} + + {% endif %} +