mirror of
https://github.com/SakanaAI/doc-to-lora.git
synced 2026-07-23 17:01:04 +02:00
webui works with 2 loras (not yet persistent)
This commit is contained in:
51616
parent
670809f555
commit
52e36fccef
6 changed files with 578 additions and 54 deletions
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
126
webui/app.py
126
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
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
</style>
|
||||
{% endblock %}
|
||||
|
||||
|
|
@ -431,9 +577,25 @@
|
|||
{% 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>
|
||||
<div id="contexts-container">
|
||||
<div class="context-header">
|
||||
<h4>Context Inputs</h4>
|
||||
<p class="instruction-text">Add one or more context inputs to generate corresponding LoRAs.</p>
|
||||
</div>
|
||||
|
||||
<div class="context-controls">
|
||||
<button id="add-context-button" class="secondary-button">Add Context</button>
|
||||
<button id="remove-context-button" class="tertiary-button" disabled>Remove Last</button>
|
||||
</div>
|
||||
|
||||
<div class="context-fields">
|
||||
<div class="context-field">
|
||||
<label for="context-input-0">Context 1:</label>
|
||||
<textarea id="context-input-0" class="context-input"
|
||||
placeholder="Enter context information here. This will be used by the hypernetwork to generate a LoRA for the base model."></textarea>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<button id="apply-hypernetwork-button">Apply Hypernetwork</button>
|
||||
<div id="hypernetwork-status"></div>
|
||||
|
|
@ -452,7 +614,12 @@
|
|||
</div>
|
||||
|
||||
<div id="chat-output" class="chat-output"></div>
|
||||
<input type="text" id="chat-input" class="chat-input" placeholder="Enter your message..." disabled>
|
||||
|
||||
<div class="chat-input-container">
|
||||
<textarea id="chat-input" class="chat-input textarea"
|
||||
placeholder="Enter your message... (Shift+Enter for new line)" disabled></textarea>
|
||||
<button id="send-message-button" class="primary-button" disabled>Send</button>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
|
@ -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 = `
|
||||
<label for="context-input-${newFieldIndex}">Context ${contextCount}:</label>
|
||||
<textarea id="context-input-${newFieldIndex}" class="context-input"
|
||||
placeholder="Enter additional context information here."></textarea>
|
||||
`;
|
||||
|
||||
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();
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
});
|
||||
|
||||
</script>
|
||||
|
||||
{% endblock %}
|
||||
Loading…
Add table
Add a link
Reference in a new issue