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_with_multi_loras( 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)