diff --git a/intx_sft.py b/intx_sft.py index 04c323b..1cebfd6 100755 --- a/intx_sft.py +++ b/intx_sft.py @@ -351,7 +351,7 @@ def main(): f"Loading from checkpoint: {ctx_args.from_pretrained_checkpoint}" ) model = ModulatedPretrainedModel.from_state_dict( - torch.load(ctx_args.from_pretrained_checkpoint), + torch.load(ctx_args.from_pretrained_checkpoint, weights_only=False), train=True, use_flash_attn=model_args.use_flash_attn, ) diff --git a/src/ctx_to_lora/eval.py b/src/ctx_to_lora/eval.py index fa72d7f..da0821c 100644 --- a/src/ctx_to_lora/eval.py +++ b/src/ctx_to_lora/eval.py @@ -339,7 +339,7 @@ def evaluate( assert split in ["validation", "test"] ctx_name = None if model_name_or_path is None: - state_dict = torch.load(checkpoint_path) + state_dict = torch.load(checkpoint_path, weights_only=False) ctx_name = state_dict["ctx_encoder_args"].ctx_encoder_model_name_or_path model = ModulatedPretrainedModel.from_state_dict( state_dict, diff --git a/src/ctx_to_lora/modeling_utils.py b/src/ctx_to_lora/modeling_utils.py index 8449aed..14ec4bc 100644 --- a/src/ctx_to_lora/modeling_utils.py +++ b/src/ctx_to_lora/modeling_utils.py @@ -19,6 +19,7 @@ from peft import ( LoraConfig, PeftConfig, PeftModel, + PeftMixedModel, LoraRuntimeConfig, get_peft_model_state_dict, set_peft_model_state_dict, @@ -52,6 +53,7 @@ from ctx_to_lora.hooks import ( from ctx_to_lora.model_loading import get_lora_config, get_model, get_model_and_tokenizer from ctx_to_lora.pooling import POOL_FN, get_pooling_fn from ctx_to_lora.utils import ( + generated_lora_to_state_dict, get_lora_module_names, get_num_layers, get_peft_in_out_features, @@ -1079,7 +1081,7 @@ class HyperLoRA(nn.Module): class ModulatedPretrainedModel(nn.Module): def __init__( self, - base_model: PreTrainedModel, + base_model: PeftModel, hypernet_config: HypernetConfig, ctx_encoder_args: CtxEncoderArguments, # use_kl_loss is only used for training @@ -1092,6 +1094,8 @@ class ModulatedPretrainedModel(nn.Module): self.ctx_encoder_args = ctx_encoder_args self.use_kl_loss = use_kl_loss self.use_base_input_as_ctx = use_base_input_as_ctx + self.lora_id = 1 + self.active_adapters = [] self.register_module("base_model", base_model) self._init_model() @@ -1107,6 +1111,7 @@ class ModulatedPretrainedModel(nn.Module): use_flash_attn: bool = True, ): lora_config = state_dict["hypernet_config"].lora_config + print(f"lora_config: {lora_config}") model_name_or_path = state_dict["base_model_name_or_path"] base_model = get_model( model_name_or_path, @@ -1471,9 +1476,115 @@ class ModulatedPretrainedModel(nn.Module): else: return model_outputs + # @contextmanager + @torch.inference_mode() + def generate_with_multi_loras( + self, + ctx_ids: Integer[Tensor, "bs ctx_length"], + ctx_attn_mask: Optional[Integer[Tensor, "bs ctx_length"]] = None, + ctx_position_ids: Optional[Integer[Tensor, "bs ctx_length"]] = None, + *model_inputs_args: Any, + **model_inputs_kwargs: dict[str, Any], + ): + # TODO: generate LoRA for each sample in ctx_ids + # apply all LoRAs to the base model + # the lora should be applied until removed + generated_loras, _ = self.generate_weights( + ctx_ids, ctx_attn_mask, ctx_position_ids + ) + # sum loras + for lora_at_module in generated_loras.values(): + if len(lora_at_module["A"]) > 1: + print("Multiple LoRAs are generated, summing them up") + lora_at_module["A"] = torch.mean(lora_at_module["A"], dim=0, keepdim=True) + lora_at_module["B"] = torch.mean(lora_at_module["B"], dim=0, keepdim=True) + + position_ids = ( + model_inputs_kwargs["position_ids"] + if "position_ids" in model_inputs_kwargs + else None + ) + + # apply loras + with apply_generated_loras( + self.base_model, + generated_loras, + self.hypernet.layer_indices, + position_ids, + self.training, + ): + model_outputs = self.base_model.generate( + *model_inputs_args, **model_inputs_kwargs + ) + return model_outputs + + # try: + # # dict of {module: + # # {A: [bs, n_layers, r, d_inim], + # # B: [bs, n_layers, r, d_outim]}} + # generated_loras, _ = self.generate_weights( + # ctx_ids, ctx_attn_mask, ctx_position_ids + # ) + # module_names = get_lora_module_names( + # self.base_model, + # self.hypernet.target_modules, + # self.hypernet.layer_indices, + # ) + + # loras = [dict() for _ in range(len(ctx_ids))] + # # unpacking the dict into a list of dicts + # # loras: list of dict {module: {A: Tensor, B: Tensor}} + # for i in range(len(ctx_ids)): + # for target_module, lora_at_module in generated_loras.items(): + # loras[i][target_module] = dict( + # A=lora_at_module["A"][i], B=lora_at_module["B"][i] + # ) + # # convert to lora state dict + # loras_sd = [ + # generated_lora_to_state_dict( + # lora, + # module_names, + # self.hypernet.target_modules, + # self.hypernet.layer_indices, + # ) + # for lora in loras + # ] + # peft_config = self.base_model.peft_config + # base_model = self.base_model + # # mixed_model = PeftMixedModel(base_model, lora_config) + # for lora_sd in loras_sd: + # lora_name = f"lora_{self.lora_id}" + # print(f"lora_name: {lora_name}") + + # # keys = list(lora_sd.keys()) + # # for key in keys: + # # new_key = key.replace("model.model", "model.model.model") + # # new_key = key.replace("weight", f"{lora_name}.weight") + # # lora_sd[new_key] = lora_sd.pop(key) + + # print(f"lora_sd: {lora_sd.keys()}") + # if lora_name not in peft_config: + # base_model.add_adapter(lora_name, peft_config["default"]) + # load_result = set_peft_model_state_dict( + # base_model, lora_sd, adapter_name=lora_name + # ) + # print(f"missing_keys: {load_result.missing_keys}") + # print(f"unexpected_keys: {load_result.unexpected_keys}") + # self.lora_id += 1 + # self.active_adapters.append(lora_name) + # base_model.set_adapter(self.active_adapters) + # print(f"active_adapters: {self.active_adapters}") + # yield base_model.base_model + + # finally: + # self.base_model.set_adapter("default") + # self.active_adapters = [] + # # self.lora_id = 1 + @torch.inference_mode() def generate( self, + # TODO: allow more than one LoRA per sample ctx_ids: Optional[Integer[Tensor, "bs ctx_length"]] = None, ctx_attn_mask: Optional[Integer[Tensor, "bs ctx_length"]] = None, ctx_position_ids: Optional[Integer[Tensor, "bs ctx_length"]] = None, @@ -1482,7 +1593,11 @@ class ModulatedPretrainedModel(nn.Module): ): generated_loras = None generated_layernorms = None - if ctx_ids is None and not self.use_base_input_as_ctx: + if ( + ctx_ids is None + and not self.active_adapters + and not self.use_base_input_as_ctx + ): logger.warning( ( "*" * 100, @@ -1493,6 +1608,14 @@ class ModulatedPretrainedModel(nn.Module): # model_outputs = self.base_model(**model_inputs_kwargs) # model_outputs.generated_loras = None # return model_outputs + if ctx_ids is None and self.active_adapters: + logger.info( + ( + "*" * 100, + "\n\nUsing active LoRAs for generation\n\n", + "*" * 100, + ) + ) else: if self.use_base_input_as_ctx: ctx_ids = ( @@ -1549,6 +1672,11 @@ class ModulatedPretrainedModel(nn.Module): return model_outputs +# def combine_loras(loras: list[LoRA]) -> LoRA: +# # TODO: implement +# ... + + class ModulatedModelWithSharedInput(nn.Module): def __init__( self, @@ -1681,7 +1809,6 @@ def apply_generated_loras( hooks = [] for module_name in generated_loras: for layer_idx in layer_indices: - # TODO: handle sequence packing??? hooks += add_generated_lora_hook( base_model, module_name, @@ -1786,6 +1913,7 @@ if __name__ == "__main__": state_dict = torch.load( "train_outputs/runs/Jan15_19-10-21_slurm0-a3nodeset-11_d1842c41/checkpoint-55000/pytorch_model.bin", + weights_only=False, ) print(state_dict.keys()) breakpoint() diff --git a/src/ctx_to_lora/run_lm_eval.py b/src/ctx_to_lora/run_lm_eval.py index a04ad5e..be9153b 100644 --- a/src/ctx_to_lora/run_lm_eval.py +++ b/src/ctx_to_lora/run_lm_eval.py @@ -75,7 +75,7 @@ if __name__ == "__main__": ) if "pytorch_model.bin" in inp: checkpoint_path = inp - state_dict = torch.load(checkpoint_path) + state_dict = torch.load(checkpoint_path, weights_only=False) print(f"Evaluating {checkpoint_path}") if "checkpoint" in checkpoint_path: diff --git a/webui/app.py b/webui/app.py index d372b57..93f2071 100644 --- a/webui/app.py +++ b/webui/app.py @@ -331,27 +331,37 @@ def load_checkpoint(): run = request.form["run"] checkpoint = request.form["checkpoint"] - context = request.form.get("context", "") + contexts = request.form.getlist("contexts[]") + + # Filter out empty contexts + contexts = [ctx for ctx in contexts if ctx.strip()] + + # If no contexts provided, use a single empty string + if not contexts: + contexts = [""] + + print(f"Received {len(contexts)} non-empty contexts for LoRA generation") 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) + state_dict = torch.load(checkpoint_path, weights_only=False) modulated_model = ModulatedPretrainedModel.from_state_dict( state_dict, train=False, - use_flash_attn=False, + use_flash_attn=True, ) modulated_model = modulated_model.to(device) modulated_model.eval() - result = {"success": True, "message": f"Loaded checkpoint {checkpoint}"} - - if context: - result["context_processed"] = True + result = { + "success": True, + "message": f"Loaded checkpoint {checkpoint}", + "contexts_processed": len(contexts), + } return jsonify(result) except Exception as e: @@ -362,6 +372,49 @@ def load_checkpoint(): return jsonify({"success": False, "error": str(e)}) +def process_multiple_contexts(contexts, ctx_tokenizer): + """ + Process multiple contexts for the ModulatedPretrainedModel. + + Args: + contexts: List of context strings + ctx_tokenizer: The tokenizer for the context model + max_length: Maximum length for tokenization + + Returns: + Dictionary with ctx_ids and ctx_attn_mask as tensors + """ + # Remove empty contexts + contexts = [ctx for ctx in contexts if ctx.strip()] + + # If no contexts provided, use a single empty string + if not contexts: + contexts = [""] + + print(f"Processing {len(contexts)} non-empty contexts") + + # Tokenize each context + all_ctx_ids = [] + all_ctx_attn_mask = [] + + for context in contexts: + # Tokenize the single context + inputs = ctx_tokenizer(context, return_tensors="pt") + + all_ctx_ids.append(inputs["input_ids"][0]) + all_ctx_attn_mask.append(inputs["attention_mask"][0]) + + # Pad to the same length + ctx_ids = torch.nn.utils.rnn.pad_sequence( + all_ctx_ids, batch_first=True, padding_value=ctx_tokenizer.pad_token_id + ) + ctx_attn_mask = torch.nn.utils.rnn.pad_sequence( + all_ctx_attn_mask, batch_first=True, padding_value=0 + ) + + return {"ctx_ids": ctx_ids, "ctx_attn_mask": ctx_attn_mask} + + @app.route("/chat", methods=["POST"]) def chat(): """ @@ -388,8 +441,17 @@ def chat(): with torch.inference_mode(), torch.amp.autocast(str(device)): from transformers import AutoTokenizer - # Get the context and tokenize it - context = request.form.get("context", "") + # Get the contexts and tokenize them + contexts = request.form.getlist("contexts[]") + if not contexts: + contexts = [""] # Use empty context if none provided + + # Ensure we have at least one context + if len(contexts) == 1 and not contexts[0].strip(): + contexts = [""] + + print(f"Processing {len(contexts)} contexts for response generation") + print(f"Contexts: {contexts}") 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 @@ -402,22 +464,20 @@ def chat(): 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]), + # Process the contexts for the ctx_encoder + ctx_inputs = process_multiple_contexts( + contexts, ctx_tokenizer, + # max_length=( + # modulated_model.ctx_encoder_args.max_ctx_len + # if hasattr(modulated_model.ctx_encoder_args, "max_ctx_len") + # else 512 + # ), ) - # 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 - ) + + ctx_ids = ctx_inputs["ctx_ids"].to(device) + ctx_attn_mask = ctx_inputs["ctx_attn_mask"].to(device) print(f"chat_history: {chat_history}") @@ -427,14 +487,32 @@ def chat(): ).to(device) # Generate response with context-modulated model - outputs = modulated_model.generate( - input_ids=model_inputs, + # with modulated_model.generate_and_apply_loras( + # ctx_ids=ctx_ids, + # ctx_attn_mask=ctx_attn_mask, + # ) as applied_model: + # outputs = applied_model.generate( + # input_ids=model_inputs, + # max_new_tokens=512, + # do_sample=False, + # ) + + outputs = modulated_model.generate_with_multi_loras( ctx_ids=ctx_ids, ctx_attn_mask=ctx_attn_mask, + input_ids=model_inputs, max_new_tokens=512, do_sample=False, ) + # outputs = modulated_model.generate( + # ctx_ids=ctx_ids, + # ctx_attn_mask=ctx_attn_mask, + # input_ids=model_inputs, + # 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 diff --git a/webui/templates/visualize.html b/webui/templates/visualize.html index 0688595..d326b6e 100644 --- a/webui/templates/visualize.html +++ b/webui/templates/visualize.html @@ -179,14 +179,62 @@ .chat-container { margin-bottom: 20px; + border: 1px solid #ddd; + border-radius: 10px; + padding: 15px; + background-color: #fafafa; + box-shadow: 0 2px 4px rgba(0, 0, 0, 0.1); + } + + .chat-container h2 { + margin-top: 0; + margin-bottom: 15px; + color: #333; + border-bottom: 1px solid #eee; + padding-bottom: 10px; } .chat-output { - min-height: 150px; + min-height: 200px; + max-height: 400px; border: 1px solid #ddd; + border-radius: 8px; padding: 10px; margin-bottom: 10px; overflow-y: auto; + background-color: white; + } + + .chat-output p { + margin: 8px 0; + white-space: pre-wrap; + word-wrap: break-word; + line-height: 1.4; + } + + .chat-message { + padding: 8px 12px; + border-radius: 10px; + margin-bottom: 10px; + max-width: 85%; + } + + .chat-message.user { + background-color: #e3f2fd; + margin-left: auto; + } + + .chat-message.assistant { + background-color: #f5f5f5; + margin-right: auto; + } + + .chat-message.system { + background-color: #fffde7; + font-style: italic; + width: 100%; + text-align: center; + font-size: 0.9em; } .chat-input { @@ -300,6 +348,104 @@ color: #333; font-size: 0.9em; } + + .context-field { + margin-bottom: 15px; + padding: 15px; + border: 1px solid #e0e0e0; + border-radius: 5px; + background-color: #f9f9f9; + } + + .context-field label { + margin-bottom: 8px; + color: #333; + font-weight: bold; + } + + .context-input { + width: 100%; + height: 100px; + resize: vertical; + padding: 8px; + border: 1px solid #ccc; + border-radius: 5px; + } + + .context-controls { + display: flex; + gap: 10px; + margin-bottom: 15px; + } + + .tertiary-button { + background-color: #ff5252; + color: white; + padding: 8px 16px; + border: none; + border-radius: 5px; + cursor: pointer; + transition: background-color 0.3s; + } + + .tertiary-button:hover { + background-color: #d32f2f; + } + + .tertiary-button:disabled { + background-color: #ffcdd2; + cursor: not-allowed; + } + + .context-header { + margin-bottom: 15px; + } + + .context-header h4 { + margin: 0 0 5px 0; + color: #333; + } + + .chat-input.textarea { + min-height: 60px; + max-height: 200px; + resize: vertical; + overflow-y: auto; + font-family: inherit; + font-size: inherit; + padding: 10px; + box-sizing: border-box; + } + + .chat-input-container { + display: flex; + gap: 10px; + margin-top: 10px; + } + + .chat-input-container textarea { + flex: 1; + } + + .primary-button { + background-color: #2196F3; + color: white; + padding: 8px 16px; + border: none; + border-radius: 5px; + cursor: pointer; + transition: background-color 0.3s; + align-self: flex-end; + } + + .primary-button:hover { + background-color: #0d8aee; + } + + .primary-button:disabled { + background-color: #bbdefb; + cursor: not-allowed; + } {% endblock %} @@ -431,9 +577,25 @@ {% endfor %} - - +
+
+

Context Inputs

+

Add one or more context inputs to generate corresponding LoRAs.

+
+ +
+ + +
+ +
+
+ + +
+
+
@@ -452,7 +614,12 @@
- + +
+ + +
@@ -557,14 +724,18 @@ document.getElementById('chat-input').value = ''; addMessageToChat('You: ' + message); - // Check if we're using hypernetwork and include context + // Prepare form data 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); + // Check if we're using hypernetwork and include contexts + if (document.getElementById('chat-model-name').textContent.includes('Hypernetwork')) { + const contextInputs = document.querySelectorAll('.context-input'); + + // Add all contexts to the form data + Array.from(contextInputs).forEach(input => { + formData.append('contexts[]', input.value); + }); } fetch('/chat', { @@ -602,15 +773,20 @@ if (applyHypernetworkButton) { applyHypernetworkButton.addEventListener('click', function () { const checkpoint = document.getElementById('checkpoint-select').value; - const context = document.getElementById('context-input').value; + const contextInputs = document.querySelectorAll('.context-input'); + + // Collect all contexts + const contexts = Array.from(contextInputs).map(input => 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?')) { + // Check if all contexts are empty + const allEmpty = contexts.every(context => !context.trim()); + if (allEmpty) { + if (!confirm('All context fields are empty. Are you sure you want to continue without any context?')) { return; } } @@ -623,14 +799,19 @@ applyHypernetworkButton.disabled = true; applyHypernetworkButton.textContent = 'Loading...'; + // Build the form data with all contexts + const formData = new FormData(); + formData.append('run', '{{ run }}'); + formData.append('checkpoint', checkpoint); + + // Add all contexts to the form data + contexts.forEach(context => { + formData.append('contexts[]', context); + }); + fetch('/load_checkpoint', { method: 'POST', - headers: { - 'Content-Type': 'application/x-www-form-urlencoded', - }, - body: 'run=' + encodeURIComponent('{{ run }}') + - '&checkpoint=' + encodeURIComponent(checkpoint) + - '&context=' + encodeURIComponent(context) + body: formData }) .then(response => response.json()) .then(data => { @@ -639,18 +820,31 @@ applyHypernetworkButton.textContent = 'Apply Hypernetwork'; if (data.success) { - statusEl.textContent = 'Success: ' + data.message; + const contextCount = data.contexts_processed || 0; + statusEl.textContent = `Success: ${data.message} (${contextCount} contexts processed)`; 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; + + // Enable the chat input + const chatInput = document.getElementById('chat-input'); + const sendButton = document.getElementById('send-message-button'); + if (chatInput) { + chatInput.disabled = false; + // Still disabled until text is entered + if (sendButton) sendButton.disabled = true; + } + 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.'); + if (contextCount > 1) { + addMessageToChat(`System: ${contextCount} contexts were processed to generate combined LoRAs.`); + } } else { statusEl.textContent = 'Error: ' + data.error; @@ -694,7 +888,12 @@ document.getElementById('chat-model-name').textContent = data.model_name; document.getElementById('load-chat-container').style.display = 'none'; document.getElementById('chat-container').style.display = 'block'; - document.getElementById('chat-input').disabled = false; + + // Enable the input and send button + const chatInput = document.getElementById('chat-input'); + const sendButton = document.getElementById('send-message-button'); + chatInput.disabled = false; + if (sendButton) sendButton.disabled = true; // Still disabled until text is entered // Show the hypernetwork container now that base model is loaded const hypernetworkContainer = document.getElementById('hypernetwork-container'); @@ -717,8 +916,42 @@ function addMessageToChat(message) { var chatOutput = document.getElementById('chat-output'); - var newMessage = document.createElement('p'); - newMessage.textContent = message; + var newMessage = document.createElement('div'); + + // Determine message type + if (message.startsWith('You: ')) { + newMessage.className = 'chat-message user'; + newMessage.textContent = message.replace('You: ', ''); + } else if (message.startsWith('System: ')) { + newMessage.className = 'chat-message system'; + newMessage.textContent = message.replace('System: ', ''); + } else { + newMessage.className = 'chat-message assistant'; + const colonIndex = message.indexOf(': '); + if (colonIndex !== -1) { + const sender = message.substring(0, colonIndex); + const content = message.substring(colonIndex + 2); + + // Add a small label for the model name + const modelLabel = document.createElement('div'); + modelLabel.className = 'model-label'; + modelLabel.style.fontSize = '0.8em'; + modelLabel.style.fontWeight = 'bold'; + modelLabel.style.marginBottom = '5px'; + modelLabel.style.color = '#555'; + modelLabel.textContent = sender; + + newMessage.appendChild(modelLabel); + + // Add the message content in a separate div + const messageContent = document.createElement('div'); + messageContent.textContent = content; + newMessage.appendChild(messageContent); + } else { + newMessage.textContent = message; + } + } + chatOutput.appendChild(newMessage); chatOutput.scrollTop = chatOutput.scrollHeight; } @@ -773,6 +1006,91 @@ } }); + // Add JavaScript to manage multiple context fields + document.addEventListener('DOMContentLoaded', function () { + const addContextButton = document.getElementById('add-context-button'); + const removeContextButton = document.getElementById('remove-context-button'); + const contextFieldsContainer = document.querySelector('.context-fields'); + + let contextCount = 1; // We start with one context field + + // Add a new context field + addContextButton.addEventListener('click', function () { + const newFieldIndex = contextCount; + contextCount++; + + const newField = document.createElement('div'); + newField.className = 'context-field'; + newField.innerHTML = ` + + + `; + + contextFieldsContainer.appendChild(newField); + + // Enable the remove button if we have more than one context + if (contextCount > 1) { + removeContextButton.disabled = false; + } + }); + + // Remove the last context field + removeContextButton.addEventListener('click', function () { + if (contextCount > 1) { + contextFieldsContainer.removeChild(contextFieldsContainer.lastChild); + contextCount--; + + // Disable the remove button if we're back to just one context + if (contextCount === 1) { + removeContextButton.disabled = true; + } + } + }); + }); + + // Add event handler for the send button + document.addEventListener('DOMContentLoaded', function () { + const sendButton = document.getElementById('send-message-button'); + const chatInput = document.getElementById('chat-input'); + + if (sendButton && chatInput) { + // Enable/disable the button based on chat input state + const updateSendButtonState = function () { + sendButton.disabled = chatInput.disabled || !chatInput.value.trim(); + }; + + // Initial state + updateSendButtonState(); + + // Add event listener for input changes + chatInput.addEventListener('input', updateSendButtonState); + + // Add event listener for send button click + sendButton.addEventListener('click', function () { + if (!chatInput.disabled && chatInput.value.trim()) { + sendMessage(); + } + }); + + // Update the event listener for chat input to handle Shift+Enter + chatInput.addEventListener('keydown', function (e) { + if (e.key === 'Enter') { + if (e.shiftKey) { + // Let Shift+Enter add a new line - no need to do anything + return; + } else { + // Just Enter triggers send + e.preventDefault(); + if (!chatInput.disabled && chatInput.value.trim()) { + sendMessage(); + } + } + } + }); + } + }); + {% endblock %} \ No newline at end of file