diff --git a/README.md b/README.md index 6e0ea18..d484f7b 100644 --- a/README.md +++ b/README.md @@ -1,12 +1,29 @@ +# Doc-to-LoRA (D2L): Learning to Instantly Internalize Contexts
-

Doc-to-LoRA

-
- + :newspaper:X | + :scroll:Paper | + :hugs:Hugging Face | + :octocat:GitHub +
A reference implementation of Doc-to-LoRA (D2L).
+
+
+
- --- +## ๐Ÿ› ๏ธ Installation +``` +curl -LsSf https://astral.sh/uv/install.sh | sh +./install.sh +``` + +## ๐Ÿค— Pre-Trained Models +``` +uv run huggingface-cli login +uv run huggingface-cli download SakanaAI/doc-to-lora --local-dir . --include "trained_t2l/*" +``` + ## ๐Ÿš€ Python API Usage ```python # caveat: this interface only supports non-batched inputs @@ -51,10 +68,14 @@ outputs = model.generate(input_ids=chat_ids, max_new_tokens=256) print(tokenizer.decode(outputs[0])) ``` -### ๐ŸŽฎ Interactive Demo [WIP] +### ๐ŸŽฎ Interactive Demo ```bash uv run demo/app.py ``` +
+

Video Demo

+ +
### ๐Ÿงช Experimental Scripts To run any of the following scripts, use `uv run $PATH_TO_SCRIPT` from the root of this project. @@ -66,14 +87,21 @@ To run any of the following scripts, use `uv run $PATH_TO_SCRIPT` from the root | [NIAH](scripts/niah/) | `scripts/niah/0-gen_data.sh` | `scripts/niah/1-train.sh` | `scripts/niah/2-eval.sh` | Run the scripts in order; data generation only needs to happen once | -### Self-Generated Data Viewer +### ๐Ÿ”ฌ Self-Generated Data Viewer After downloading/generating the data, we can see samples of the data using this script. ```bash uv run webui/self_gen_viewer.py ``` See more info at [webui/SELF_GEN_VIEWER.md](webui/SELF_GEN_VIEWER.md). -### Citation +### ๐Ÿ“š Citation ```bibtex -TBD -``` \ No newline at end of file +@techreport{sakana2025doc-to-lora, + title = {{Doc-to-LoRA: Learning to Instantly Internalize Contexts}}, + author = {Rujikorn Charakorn and Edoardo Cetin and Shinnosuke Uesaka and Robert Tjarko Lange}, + institution = {Sakana AI}, + year = {2026}, + month = {Febuary}, + note = {Technical Report} +} +``` diff --git a/assets/doc-to-lora-demo.mp4 b/assets/doc-to-lora-demo.mp4 new file mode 100644 index 0000000..a78146c Binary files /dev/null and b/assets/doc-to-lora-demo.mp4 differ diff --git a/assets/overview_animation.gif b/assets/overview_animation.gif new file mode 100644 index 0000000..1b075ba Binary files /dev/null and b/assets/overview_animation.gif differ diff --git a/demo/app.py b/demo/app.py index 1de714c..f9eb6e6 100644 --- a/demo/app.py +++ b/demo/app.py @@ -145,12 +145,9 @@ def load_checkpoint( ) -def process_multiple_contexts(contexts: list[str]) -> dict: - contexts = [ctx for ctx in contexts if ctx.strip()] - if not contexts: - contexts = [""] - print(f"Processing {len(contexts)} non-empty contexts") - tokenized_contexts = tokenize_ctx_text({"context": contexts}, ctx_tokenizer) +def process_context(context: str) -> dict: + context = context.strip() if context else "" + tokenized_contexts = tokenize_ctx_text({"context": [context]}, ctx_tokenizer) ctx_ids = tokenized_contexts["ctx_ids"] ctx_ids = [ torch.tensor(ctx_id, dtype=torch.long, device=device) for ctx_id in ctx_ids @@ -181,14 +178,8 @@ def add_user_message(message: str, history): def generate_response( history: list[list[str]], system_msg: str, - context_1: str, - scaler_1: float, - context_2: str, - scaler_2: float, - context_3: str, - scaler_3: float, - context_4: str, - scaler_4: float, + context: str, + context_scaler: float, bias_scaler: float, ): global modulated_model, chat_history, ctx_tokenizer, base_tokenizer @@ -210,41 +201,24 @@ def generate_response( chat_history.append({"role": "user", "content": user_message}) - contexts = [] - scalers = [] - - if context_1 and context_1.strip(): - contexts.append(context_1) - scalers.append(scaler_1) - if context_2 and context_2.strip(): - contexts.append(context_2) - scalers.append(scaler_2) - if context_3 and context_3.strip(): - contexts.append(context_3) - scalers.append(scaler_3) - if context_4 and context_4.strip(): - contexts.append(context_4) - scalers.append(scaler_4) - - if not contexts: - contexts = [""] - scalers = [1.0] - - print(f"Processing {len(contexts)} contexts with scalers: {scalers}") + context = context.strip() if context else "" + print(f"Processing single context with scaler: {context_scaler}") print(f"Bias scaler: {bias_scaler}") with torch.inference_mode(), torch.amp.autocast(str(device)): - ctx_inputs = process_multiple_contexts(contexts) + ctx_inputs = process_context(context) ctx_ids = ctx_inputs["ctx_ids"].to(device) ctx_attn_mask = ctx_inputs["ctx_attn_mask"].to(device) - scalers_tensor = torch.tensor(scalers, dtype=torch.float32, device=device) + scalers_tensor = torch.tensor( + [context_scaler], dtype=torch.float32, device=device + ) model_inputs = base_tokenizer.apply_chat_template( chat_history, return_tensors="pt", add_generation_prompt=True ).to(device) - print(f"Contexts: {contexts}") + print(f"Context: {context}") print(f"Chat history: {chat_history}") outputs = modulated_model.generate( @@ -292,13 +266,6 @@ def reset_chat(system_msg: str): return [[None, WARNING_MESSAGE]], "Chat history reset successfully!" -def update_context_display(num_contexts: int): - updates = [] - for i in range(4): - updates.append(gr.update(visible=(i < num_contexts))) - return updates - - custom_css = """ :root { color-scheme: light; @@ -492,7 +459,7 @@ def create_demo(): # ๐Ÿ“œ Doc-to-LoRA Chat Interface Load a hypernetwork checkpoint and chat with a context-modulated language model. - Add multiple contexts with individual scaling parameters to influence the model's responses. + Add one context with a scaling parameter to influence the model's responses. """ ) @@ -512,10 +479,9 @@ def create_demo(): ### ๐Ÿ“– Usage Instructions 1. **Load a Checkpoint**: Select a hypernetwork checkpoint from the dropdown and click "Load Checkpoint" - 2. **Configure Contexts**: - - Use the slider to choose how many contexts you want (1-4) - - Enter your context information in each text field - - Adjust the scaling sliders to control the influence of each context + 2. **Configure Context**: + - Enter your context information in the text field + - Adjust the scaling slider to control context influence 3. **Set Bias Scaler**: Adjust the bias scaler to control overall model behavior 4. **Start Chatting**: Once the model is loaded, type your message and press Shift+Enter or click Send 5. **Reset**: Use the "Reset Chat" button to start a new conversation @@ -555,76 +521,24 @@ def create_demo(): """
๐Ÿง  Context Internalization (Hypernetwork Input)
- These contexts modulate the model internally โ€” they are NOT shown to the base model + This context modulates the model internally โ€” it is NOT shown to the base model
""" ) - num_contexts = gr.Slider( - minimum=1, - maximum=4, - step=1, - value=1, - label="Number of Contexts to Internalize", - interactive=True, + context = gr.Textbox( + label="๐Ÿง  Context (Internalized via Hypernetwork)", + placeholder="Enter context to be internalized by the hypernetwork...", + lines=4, + value=DEFAULT_CONTEXT, + ) + context_scaler = gr.Slider( + minimum=-2.0, + maximum=2.0, + step=0.01, + value=1.0, + label="Context Scaling", ) - - with gr.Group(visible=True) as context_1_group: - context_1 = gr.Textbox( - label="๐Ÿง  Context 1 (Internalized via Hypernetwork)", - placeholder="Enter context to be internalized by the hypernetwork...", - lines=4, - value=DEFAULT_CONTEXT, - ) - scaler_1 = gr.Slider( - minimum=-2.0, - maximum=2.0, - step=0.01, - value=1.0, - label="Context 1 Scaling", - ) - - with gr.Group(visible=False) as context_2_group: - context_2 = gr.Textbox( - label="๐Ÿง  Context 2 (Internalized via Hypernetwork)", - placeholder="Enter additional context to be internalized...", - lines=4, - ) - scaler_2 = gr.Slider( - minimum=-2.0, - maximum=2.0, - step=0.01, - value=1.0, - label="Context 2 Scaling", - ) - - with gr.Group(visible=False) as context_3_group: - context_3 = gr.Textbox( - label="๐Ÿง  Context 3 (Internalized via Hypernetwork)", - placeholder="Enter additional context to be internalized...", - lines=4, - ) - scaler_3 = gr.Slider( - minimum=-2.0, - maximum=2.0, - step=0.01, - value=1.0, - label="Context 3 Scaling", - ) - - with gr.Group(visible=False) as context_4_group: - context_4 = gr.Textbox( - label="๐Ÿง  Context 4 (Internalized via Hypernetwork)", - placeholder="Enter additional context to be internalized...", - lines=4, - ) - scaler_4 = gr.Slider( - minimum=-2.0, - maximum=2.0, - step=0.01, - value=1.0, - label="Context 4 Scaling", - ) gr.Markdown("---") @@ -717,17 +631,6 @@ def create_demo(): outputs=[chat_status_notice, msg, send_btn, clear_btn, chatbot], ) - num_contexts.change( - fn=update_context_display, - inputs=[num_contexts], - outputs=[ - context_1_group, - context_2_group, - context_3_group, - context_4_group, - ], - ) - msg.submit( fn=add_user_message, inputs=[msg, chatbot], @@ -737,14 +640,8 @@ def create_demo(): inputs=[ chatbot, system_msg, - context_1, - scaler_1, - context_2, - scaler_2, - context_3, - scaler_3, - context_4, - scaler_4, + context, + context_scaler, bias_scaler, ], outputs=[chatbot], @@ -759,14 +656,8 @@ def create_demo(): inputs=[ chatbot, system_msg, - context_1, - scaler_1, - context_2, - scaler_2, - context_3, - scaler_3, - context_4, - scaler_4, + context, + context_scaler, bias_scaler, ], outputs=[chatbot], diff --git a/install.sh b/install.sh index 8b1e128..2f52443 100755 --- a/install.sh +++ b/install.sh @@ -1,7 +1,3 @@ -# the installation script assumes the following modules (slurm) -# module load cuda/12.4 -# module load cudnn/9.1.0 -# module load hpcx/v2.21 curl -LsSf https://astral.sh/uv/install.sh | sh uv self update uv venv --python 3.10 --seed diff --git a/webui/app.py b/webui/app.py deleted file mode 100644 index aa9c19e..0000000 --- a/webui/app.py +++ /dev/null @@ -1,822 +0,0 @@ -import json -import os -import sys -from glob import glob - -import torch -import yaml -from flask import Flask, abort, jsonify, render_template, request -from transformers import pipeline - -from ctx_to_lora.data.processing import tokenize_ctx_text -from ctx_to_lora.model_loading import get_tokenizer -from ctx_to_lora.modeling import hypernet - -sys.modules["ctx_to_lora.modeling_utils"] = hypernet -app = Flask(__name__) - -TRAIN_OUTPUTS_DIR = "train_outputs/runs" -chat_generator = None -chat_model_name = None -chat_history = None -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -modulated_model = None - - -# TODO: grab the results on-deman -# TODO: cache the grabed results - - -def get_run_data(run_path): - """ - Loads and processes data from all_results.json and *_results.json for a given run. - Also recursively searches for results files in subfolders. - - Args: - run_path: Path to the directory containing the results files. - - Returns: - A dictionary containing the grouped and processed results data. - """ - results_file = os.path.join(run_path, "all_results.json") - try: - with open(results_file) as f: - data = json.load(f) - except FileNotFoundError: - data = {} - - # Load data from *_results.json files in the current directory - for results_json_file in glob(os.path.join(run_path, "*_results.json")): - try: - with open(results_json_file) as f: - split_data = json.load(f) - # Use the filename as a prefix for the keys - file_prefix = os.path.basename(results_json_file).replace( - "_results.json", "" - ) - for key, value in split_data.items(): - if isinstance(value, float): - value = round(value, 4) - new_key = f"{file_prefix}_{key}" - data[new_key] = value - except FileNotFoundError: - print(f"Warning: {results_json_file} not found.") - - # Recursively search for *_results.json files in subdirectories - for subdir in [d for d in glob(os.path.join(run_path, "*")) if os.path.isdir(d)]: - subdir_name = os.path.basename(subdir) - # Look for *_results.json files in the subdirectory - for results_json_file in glob(os.path.join(subdir, "*_results.json")): - try: - with open(results_json_file) as f: - split_data = json.load(f) - # Use a prefix that includes the subfolder name - file_prefix = f"{subdir_name}_{os.path.basename(results_json_file).replace('_results.json', '')}" - for key, value in split_data.items(): - if isinstance(value, float): - value = round(value, 4) - new_key = f"{file_prefix}_{key}" - data[new_key] = value - except FileNotFoundError: - print(f"Warning: {results_json_file} not found.") - - grouped_data = {} - for key, value in data.items(): - if isinstance(value, dict): - # Handle nested dictionaries - for subkey, subvalue in value.items(): - if isinstance(subvalue, float): - subvalue = round(subvalue, 4) - - # Determine prefix based on whether the key is from a split file or all_results - if "_" in key: - prefix = key.split("_")[0] - else: - prefix = "all" - - if prefix not in grouped_data: - grouped_data[prefix] = {} - if key not in grouped_data[prefix]: - grouped_data[prefix][key] = {} - - grouped_data[prefix][key][subkey] = subvalue - else: - if isinstance(value, float): - value = round(value, 4) - - # Determine prefix based on whether the key is from a split file or all_results - if "_" in key: - prefix = key.split("_")[0] - else: - prefix = "all" - - if prefix not in grouped_data: - grouped_data[prefix] = {} - - grouped_data[prefix][key] = value - - return grouped_data - - -def get_generated_text_data(run_path: str) -> dict: - """ - Loads generated text data from all *_generated_text.jsonl files from the current run's output. - It specifically excludes any files ending with _no_context_generated_text.jsonl, - as those are handled by get_base_model_generated_text. - Also recursively searches for generated text files in subfolders. - - Args: - run_path: Path to the directory containing the .jsonl files for the current run. - - Returns: - A dictionary containing the generated text data for each primary split found. - """ - generated_data = {} - # Find all *_generated_text.jsonl files in the current directory, excluding _no_context_ variants - files = sorted( - f - for f in glob(os.path.join(run_path, "*_generated_text.jsonl")) - if "_no_context_generated_text.jsonl" not in os.path.basename(f) - ) - print(f"Processing generated text files from {run_path}: {files}") - - for filename in files: - split = os.path.basename(filename).replace("_generated_text.jsonl", "") - try: - with open(filename) as f: - lines = f.readlines(10_000_000) - data = [json.loads(line) for line in lines] - generated_data[split] = data - except FileNotFoundError: - generated_data[split] = None - except json.JSONDecodeError as e: - print(f"Error decoding JSON from {filename}: {e}") - generated_data[split] = None - - # Recursively search for *_generated_text.jsonl files in subdirectories - for subdir_path in [ - d for d in glob(os.path.join(run_path, "*")) if os.path.isdir(d) - ]: - subdir_name = os.path.basename(subdir_path) - subdir_files = sorted( - f - for f in glob(os.path.join(subdir_path, "*_generated_text.jsonl")) - if "_no_context_generated_text.jsonl" not in os.path.basename(f) - ) - print( - f"Processing generated text files from subdirectory {subdir_path}: {subdir_files}" - ) - for filename in subdir_files: - base_split = os.path.basename(filename).replace("_generated_text.jsonl", "") - # Create a unique split key including the subdirectory name - split = f"{subdir_name}/{base_split}" - try: - with open(filename) as f: - lines = f.readlines(10_000_000) - data = [json.loads(line) for line in lines] - generated_data[split] = data - except FileNotFoundError: - generated_data[split] = None - except json.JSONDecodeError as e: - print(f"Error decoding JSON from {filename}: {e}") - generated_data[split] = None - return generated_data - - -def get_base_model_generated_text(model_name, generated_data): - """ - Loads generated text data from the base model for comparison. - Also loads "no context" base model data if available. - - Args: - model_name: The name of the base model - generated_data: Dictionary of generated data from the fine-tuned model - - Returns: - Tuple: (base_model_data, base_model_no_context_data) - Dictionaries containing base model generated texts for matching splits - """ - if not model_name: - return {}, {} - - base_model_data = {} - base_model_no_context_data = {} - - # # Create normalized model name for directory lookup - # normalized_model_name = model_name.replace("/", "_") - - for split in generated_data: - # Skip if no data for this split - if not generated_data[split]: - base_model_data[split] = None - base_model_no_context_data[split] = None - continue - - # Construct path to the base model's output for this split - base_model_path = f"eval_results/{model_name}/{split}_generated_text.jsonl" - base_model_no_context_path = ( - f"eval_results/{model_name}/{split}_no_context_generated_text.jsonl" - ) - - try: - with open(base_model_path) as f: - lines = f.readlines(10_000_000) - data = [json.loads(line) for line in lines] - base_model_data[split] = data - except FileNotFoundError: - print(f"Base model output not found for {split} at {base_model_path}") - # Store None to indicate we tried but didn't find matching data - base_model_data[split] = None - - try: - with open(base_model_no_context_path) as f: - lines = f.readlines(10_000_000) - data = [json.loads(line) for line in lines] - base_model_no_context_data[split] = data - except FileNotFoundError: - print( - f"Base model (no context) output not found for {split} at {base_model_no_context_path}" - ) - base_model_no_context_data[split] = None - - return base_model_data, base_model_no_context_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) 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) 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(): - """ - Displays a dropdown list of training runs. - """ - runs = [ - d - for d in os.listdir(TRAIN_OUTPUTS_DIR) - if os.path.isdir(os.path.join(TRAIN_OUTPUTS_DIR, d)) - ] - runs.sort( - key=lambda d: os.path.getctime(os.path.join(TRAIN_OUTPUTS_DIR, d)), - reverse=True, - ) - return render_template("index.html", runs=runs) - - -@app.route("/visualize/") -def visualize(run): - """ - Displays the results from all_results.json and the generated text data. - """ - logdir = os.path.join(TRAIN_OUTPUTS_DIR, run) - if not os.path.isdir(logdir): - abort(404, description=f"Run '{run}' not found.") - - eval_folders = [ - d - for d in os.listdir(logdir) - if os.path.isdir(os.path.join(logdir, d)) and d.startswith("eval-results-") - ] - eval_folders.sort( - key=lambda d: os.path.getctime(os.path.join(logdir, d)), - 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: - # New structure with eval folders and a folder is selected - eval_path = os.path.join(logdir, selected_eval_folder) - data = get_run_data(eval_path) - if data is None: - abort( - 404, - description=f"'all_results.json' not found in '{selected_eval_folder}'.", - ) - generated_data = get_generated_text_data(eval_path) - elif not eval_folders: - # Old structure: eval results in the root folder - data = get_run_data(logdir) - if data is None: - abort(404, description=f"'all_results.json' not found in '{run}'.") - generated_data = get_generated_text_data(logdir) - selected_eval_folder = "root" # Indicate root folder for template - else: - # New structure, but no eval folder selected yet - data = None - generated_data = None - - # Load config.yaml from the run directory (root level) - 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: - model_name = None - - # If config.yaml doesn't have the model name, try args.yaml - if not model_name: - args_path = os.path.join(logdir, "args.yaml") - try: - with open(args_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, - ) - args = yaml.safe_load(f) - model_name = args.get("model_name_or_path") - except FileNotFoundError: - model_name = None - - # Get base model outputs if we have the model name and modulated outputs - base_model_data = {} - base_model_no_context_data = {} - if model_name and generated_data: - base_model_data, base_model_no_context_data = get_base_model_generated_text( - model_name, generated_data - ) - - return render_template( - "visualize.html", - run=run, - eval_folders=eval_folders, - selected_eval_folder=selected_eval_folder, - data=data, - generated_data=generated_data, - base_model_data=base_model_data, - base_model_no_context_data=base_model_no_context_data, - model_name=model_name, - checkpoints=checkpoint_names, - ) - - -@app.route("/load_model", methods=["POST"]) -def load_model(): - """ - Loads the model and creates the pipeline. - """ - 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, "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, tokenizer=tokenizer, device=device - ) - return jsonify({"model_name": chat_model_name}) - else: - 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.hypernet import ModulatedPretrainedModel - - run = request.form["run"] - checkpoint = request.form["checkpoint"] - 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, weights_only=False) - - modulated_model = ModulatedPretrainedModel.from_state_dict( - state_dict, - train=False, - use_flash_attn=True, - use_sequence_packing=False, # for generation - ) - modulated_model = modulated_model.to(device).to(torch.bfloat16) - modulated_model.eval() - - result = { - "success": True, - "message": f"Loaded checkpoint {checkpoint}", - "contexts_processed": len(contexts), - } - - 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)}) - - -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") - - tokenized_contexts = tokenize_ctx_text({"context": contexts}, ctx_tokenizer) - ctx_ids = tokenized_contexts["ctx_ids"] - ctx_ids = [ - torch.tensor(ctx_id, dtype=torch.long, device=device) for ctx_id in ctx_ids - ] - ctx_attn_mask = [torch.ones_like(ids) for ids in ctx_ids] - ctx_attn_mask = [ - torch.tensor(ctx_attn_mask, dtype=torch.long, device=device) - for ctx_attn_mask in ctx_attn_mask - ] - ctx_ids = torch.nn.utils.rnn.pad_sequence( - ctx_ids, - batch_first=True, - padding_value=0, - ) - ctx_attn_mask = torch.nn.utils.rnn.pad_sequence( - 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(): - """ - 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 or modulated_model: - print("Generating 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)): - # Get the contexts and tokenize them - raw_contexts = request.form.getlist("contexts[]") - # Build scalers aligned with contexts (default to 1.0) - raw_scalers = request.form.getlist("scalers[]") - # Parse bias scaler (singular), default to 1.0 if missing/invalid - bias_scaler_str = request.form.get("bias_scaler", "1.0") - try: - bias_scaler = float(bias_scaler_str) - except Exception: - bias_scaler = 1.0 - - # Clean contexts (drop empty), mirror selection for scalers - pairs = list(zip(raw_contexts, raw_scalers)) - contexts = [] - kept_scalers = [] - for ctx, sc in pairs: - if ctx.strip(): - contexts.append(ctx) - try: - kept_scalers.append(float(sc)) - except Exception: - kept_scalers.append(1.0) - - # If all contexts are empty, use a single empty context and scaler 1.0 - if not contexts: - contexts = [""] - kept_scalers = [1.0] - - # Ensure we have at least one context (safety) - if len(contexts) == 1 and not contexts[0].strip(): - contexts = [""] - if not kept_scalers: - kept_scalers = [1.0] - - 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 - ) - 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 - ) - - # Process the contexts for the ctx_encoder - ctx_inputs = process_multiple_contexts( - contexts, - ctx_tokenizer, - ) - - ctx_ids = ctx_inputs["ctx_ids"].to(device) - ctx_attn_mask = ctx_inputs["ctx_attn_mask"].to(device) - - scalers_tensor = torch.tensor( - kept_scalers, dtype=torch.float32, device=device - ) - - print(f"chat_history: {chat_history}") - print(f"scalers: {scalers_tensor}") - print(f"bias_scaler: {bias_scaler}") - - # Tokenize the chat history - model_inputs = base_tokenizer.apply_chat_template( - chat_history, return_tensors="pt", add_generation_prompt=True - ).to(device) - - # Generate response with context-modulated model - # 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( - ctx_ids=ctx_ids, - ctx_attn_mask=ctx_attn_mask, - n_ctx_chunks=torch.tensor( - [len(ctx_ids)], device=ctx_ids.device - ), - scalers=scalers_tensor, # pass per-context scalers - bias_scaler=bias_scaler, # pass singular bias scaler - 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 - ) - - 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: - return jsonify( - { - "response": "No model loaded. Please load a base model or apply a hypernetwork first." - } - ) - - -@app.route("/update_system_msg", methods=["POST"]) -def update_system_msg(): - """ - Updates the system message in chat_history. - """ - global chat_history - data = request.get_json() - new_system_msg = data.get("system_msg", "").strip() - if new_system_msg: - if chat_history and chat_history[0]["content"] != new_system_msg: - chat_history[0]["content"] = new_system_msg - return jsonify({"success": True}) - return jsonify({"success": False}) - - -@app.route("/get_system_msg", methods=["GET"]) -def get_system_msg(): - """ - Retrieves the current system message from chat_history. - """ - global chat_history - if chat_history and "content" in chat_history[0]: - return jsonify({"system_msg": chat_history[0]["content"]}) - return jsonify({"system_msg": ""}) - - -@app.route("/get_command/") -def get_command(run): - """ - Extracts the command from the debug.log file. - - Args: - run: The name of the training run. - - Returns: - A JSON response containing the command or an error message. - """ - logdir = os.path.join(TRAIN_OUTPUTS_DIR, run) - debug_log_path = os.path.join(logdir, "debug.log") - - try: - with open(debug_log_path) as f: - for line in f: - if "CMD:" in line: - command = line.split("CMD:")[1].strip() - return jsonify({"command": command}) - return jsonify({"command": "Command not found in debug.log."}) - except FileNotFoundError: - return jsonify({"command": "debug.log not found."}) - - -@app.route("/get_config/") -def get_config(run): - """ - Extracts the config file name from cli_args.yaml. - - Args: - run: The name of the training run. - - Returns: - A JSON response containing the config file name or an error message. - """ - logdir = os.path.join(TRAIN_OUTPUTS_DIR, run) - cli_args_path = os.path.join(logdir, "cli_args.yaml") - - try: - with open(cli_args_path) as f: - cli_args = yaml.safe_load(f) - config_file = cli_args.get("config", "Config file not found.") - return jsonify({"config": config_file}) - except FileNotFoundError: - 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/base.html b/webui/templates/base.html deleted file mode 100644 index ff61f94..0000000 --- a/webui/templates/base.html +++ /dev/null @@ -1,22 +0,0 @@ - - - - - {% block title %}My Web App{% endblock %} - - {% block head %}{% endblock %} - - - - {% block content %}{% endblock %} - - - \ No newline at end of file diff --git a/webui/templates/index.html b/webui/templates/index.html deleted file mode 100644 index dddb75e..0000000 --- a/webui/templates/index.html +++ /dev/null @@ -1,37 +0,0 @@ -{% extends "base.html" %} - -{% block title %}Training Run Visualizer{% endblock %} - -{% block content %} -

Training Runs

- - - - - - - - - - {% for run in runs %} - - - - - - - {% endfor %} - -
Run NameCommandConfig File
{{ run }}Loading...Loading...
-{% endblock %} \ No newline at end of file diff --git a/webui/templates/visualize.html b/webui/templates/visualize.html deleted file mode 100644 index d25045e..0000000 --- a/webui/templates/visualize.html +++ /dev/null @@ -1,1265 +0,0 @@ -{% extends "base.html" %} - -{% block title %}Training Run Visualization: {{ run }}{% endblock %} - -{% block head %} - -{% endblock %} - -{% block content %} -
-
-

Training Run: {{ run }}

-

Command:

-
- -
- -
-
- {% if eval_folders %} -
- - -
- {% endif %} - - {% if data %} - {% for group, results in data.items() %} -

{{ group }} ({% if selected_eval_folder and selected_eval_folder != "root" %}{{ selected_eval_folder }}{% - else %}root{% endif %})

- - - - - - - - - {% for metric, value in results.items() %} - - - {% if value is mapping %} - - {% else %} - - {% endif %} - - {% endfor %} - -
MetricValue
{{ metric }} - - - - - - - - - {% for submetric, subvalue in value.items() %} - - - - - {% endfor %} - -
Sub-MetricValue
{{ submetric }}{{ subvalue }}
- 		{{ value }}
- {% endfor %} - {% endif %} -
- -
- {% if generated_data %} - {% for split, data in generated_data.items() %} - {% if data %} -
-

Generated Text ({{ split }} - {% if selected_eval_folder and selected_eval_folder != "root" %}{{ - selected_eval_folder }}{% else %}root{% endif %})

-
- -

Input:

- - -
-
-

Modulated Model:

-
-
-

Base Model ({{ model_name }}): Not - available

-
-
-

Base Model (No Context - {{ model_name }}): Not available

-
-
- -

Label:

-
-
- - - / {{ data|length }} - - -
-
- {% endif %} - {% endfor %} - {% elif selected_eval_folder %} -

No generated text data found in {% if selected_eval_folder and selected_eval_folder != "root" %}{{ - 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 %} - - -
-
- - - -{% endblock %} \ No newline at end of file