webui working (not yet two loras)

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

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 %}

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}")

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)

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;
+    }
 </style>
 {% endblock %}
 
@@ -333,16 +403,52 @@
             selected_eval_folder }}{% else %}root{% endif %}.</p>
         {% endif %}
 
+        <!-- Step-by-step guide -->
+        <div class="ui-sequence">
+            <h3>How to use this interface:</h3>
+            <ol>
+                <li>First, click "Load Chat with {{ model_name }}" to initialize the base model</li>
+                <li>You can then chat directly with the base model</li>
+                <li>Alternatively, you can select a hypernetwork checkpoint to use context-informed responses</li>
+                <li>To reset the conversation at any time, click the "Reset Chat" button</li>
+            </ol>
+        </div>
+
         <div id="load-chat-container">
             <button id="load-chat-button">Load Chat with {{ model_name }}</button>
         </div>
+
+        <!-- Hypernetwork Section - Initially hidden -->
+        {% if checkpoints %}
+        <div class="hypernetwork-container" id="hypernetwork-container" style="display: none;">
+            <h3>Hypernetwork Context Integration</h3>
+            <p class="instruction-text">Apply a hypernetwork to modulate the base model with context information</p>
+            <label for="checkpoint-select">Select Checkpoint:</label>
+            <select id="checkpoint-select">
+                <option value="">-- Select Checkpoint --</option>
+                {% for checkpoint in checkpoints %}
+                <option value="{{ checkpoint }}">{{ checkpoint }}</option>
+                {% endfor %}
+            </select>
+
+            <label for="context-input">Context Information:</label>
+            <textarea id="context-input"
+                placeholder="Enter context information here. This will be used by the hypernetwork to generate a LoRA for the base model."></textarea>
+
+            <button id="apply-hypernetwork-button">Apply Hypernetwork</button>
+            <div id="hypernetwork-status"></div>
+        </div>
+        {% endif %}
+
         <div class="chat-container" id="chat-container" style="display: none;">
             <h2>Chat with <span id="chat-model-name"></span></h2>
 
             <div class="system-msg-container">
                 <label for="system-msg">System Message:</label>
                 <input type="text" id="system-msg" class="chat-input" placeholder="Enter system message..." value="">
-                <!-- Removed the Update System Message button -->
+                <div class="chat-controls">
+                    <button id="reset-chat-button" class="secondary-button">Reset Chat</button>
+                </div>
             </div>
 
             <div id="chat-output" class="chat-output"></div>
@@ -451,16 +557,24 @@
                 document.getElementById('chat-input').value = '';
                 addMessageToChat('You: ' + message);
 
+                // Check if we're using hypernetwork and include context
+                let formData = new FormData();
+                formData.append('message', message);
+
+                // Get the context text if hypernetwork is being used
+                const contextInput = document.getElementById('context-input');
+                if (document.getElementById('chat-model-name').textContent.includes('Hypernetwork') && contextInput) {
+                    formData.append('context', contextInput.value);
+                }
+
                 fetch('/chat', {
                     method: 'POST',
-                    headers: {
-                        'Content-Type': 'application/x-www-form-urlencoded',
-                    },
-                    body: 'message=' + encodeURIComponent(message)
+                    body: formData
                 })
                     .then(response => response.json())
                     .then(data => {
-                        addMessageToChat('{{ model_name }}: ' + data.response);
+                        const modelName = document.getElementById('chat-model-name').textContent || '{{ model_name }}';
+                        addMessageToChat(modelName + ': ' + data.response);
                     })
                     .catch(error => {
                         console.error("Error during chat fetch:", error);
@@ -482,6 +596,80 @@
             });
     }
 
+    // Handle hypernetwork checkpoint loading
+    document.addEventListener('DOMContentLoaded', function () {
+        const applyHypernetworkButton = document.getElementById('apply-hypernetwork-button');
+        if (applyHypernetworkButton) {
+            applyHypernetworkButton.addEventListener('click', function () {
+                const checkpoint = document.getElementById('checkpoint-select').value;
+                const context = document.getElementById('context-input').value;
+
+                if (!checkpoint) {
+                    alert('Please select a checkpoint first.');
+                    return;
+                }
+
+                if (!context.trim()) {
+                    if (!confirm('The context field is empty. Are you sure you want to continue without context?')) {
+                        return;
+                    }
+                }
+
+                const statusEl = document.getElementById('hypernetwork-status');
+                statusEl.textContent = 'Loading checkpoint...';
+                statusEl.style.color = '#1565c0'; // Blue to indicate in progress
+
+                // Disable the button while loading
+                applyHypernetworkButton.disabled = true;
+                applyHypernetworkButton.textContent = 'Loading...';
+
+                fetch('/load_checkpoint', {
+                    method: 'POST',
+                    headers: {
+                        'Content-Type': 'application/x-www-form-urlencoded',
+                    },
+                    body: 'run=' + encodeURIComponent('{{ run }}') +
+                        '&checkpoint=' + encodeURIComponent(checkpoint) +
+                        '&context=' + encodeURIComponent(context)
+                })
+                    .then(response => response.json())
+                    .then(data => {
+                        // Re-enable the button
+                        applyHypernetworkButton.disabled = false;
+                        applyHypernetworkButton.textContent = 'Apply Hypernetwork';
+
+                        if (data.success) {
+                            statusEl.textContent = 'Success: ' + data.message;
+                            statusEl.style.color = 'green';
+
+                            // Show the chat container if it's hidden
+                            document.getElementById('chat-container').style.display = 'block';
+                            document.getElementById('chat-input').disabled = false;
+                            document.getElementById('chat-model-name').textContent =
+                                'Hypernetwork: ' + checkpoint;
+
+                            // Clear any existing chat and add an instruction
+                            document.getElementById('chat-output').innerHTML = '';
+                            addMessageToChat('System: Hypernetwork applied successfully. Your messages will now be processed through the context-aware model.');
+
+                        } else {
+                            statusEl.textContent = 'Error: ' + data.error;
+                            statusEl.style.color = 'red';
+                        }
+                    })
+                    .catch(error => {
+                        // Re-enable the button on error
+                        applyHypernetworkButton.disabled = false;
+                        applyHypernetworkButton.textContent = 'Apply Hypernetwork';
+
+                        console.error('Error loading checkpoint:', error);
+                        statusEl.textContent = 'Error: ' + error.message;
+                        statusEl.style.color = 'red';
+                    });
+            });
+        }
+    });
+
     // Call initializeChat when chat is loaded
     document.getElementById('load-chat-button').addEventListener('click', function () {
         fetch('/load_model', {
@@ -508,8 +696,17 @@
                     document.getElementById('chat-container').style.display = 'block';
                     document.getElementById('chat-input').disabled = false;
 
+                    // Show the hypernetwork container now that base model is loaded
+                    const hypernetworkContainer = document.getElementById('hypernetwork-container');
+                    if (hypernetworkContainer) {
+                        hypernetworkContainer.style.display = 'block';
+                    }
+
                     console.log("Chat interface loaded successfully.");
                     initializeChat();
+
+                    // Add a message about UI flow
+                    addMessageToChat('System: Base model loaded successfully. You can now chat directly or apply a hypernetwork with context information.');
                 }
             })
             .catch(error => {
@@ -542,6 +739,40 @@
         }
     });
 
+    // Add the reset chat functionality
+    document.addEventListener('DOMContentLoaded', function () {
+        const resetChatButton = document.getElementById('reset-chat-button');
+        if (resetChatButton) {
+            resetChatButton.addEventListener('click', function () {
+                // Get the current system message
+                const systemMsg = document.getElementById('system-msg').value;
+
+                // Clear the chat output display
+                document.getElementById('chat-output').innerHTML = '';
+
+                fetch('/reset_chat', {
+                    method: 'POST',
+                    headers: {
+                        'Content-Type': 'application/x-www-form-urlencoded',
+                    },
+                    body: 'system_msg=' + encodeURIComponent(systemMsg)
+                })
+                    .then(response => response.json())
+                    .then(data => {
+                        if (data.success) {
+                            console.log('Chat history reset successfully');
+                            addMessageToChat('System: Chat history has been reset.');
+                        } else {
+                            console.error('Failed to reset chat history');
+                        }
+                    })
+                    .catch(error => {
+                        console.error('Error resetting chat history:', error);
+                    });
+            });
+        }
+    });
+
 </script>
 
 {% endblock %}