diff --git a/README.md b/README.md index 0fad415..6e0ea18 100644 --- a/README.md +++ b/README.md @@ -64,3 +64,16 @@ To run any of the following scripts, use `uv run $PATH_TO_SCRIPT` from the root | ------------------------------------ | ------------------------------------- | ----------------------------- | ---------------------------- | ----------------------------------------------------------------------------------------------------------------------------------- | | [Main experiment](scripts/main_exp/) | `scripts/main_exp/0-download_data.sh` | `scripts/main_exp/1-train.sh` | `scripts/main_exp/eval/*.sh` | Downloading data is fastest; regenerate only if you need fresh synthetic data. Evaluation scripts reproduce the main paper metrics. | | [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 +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 +```bibtex +TBD +``` \ No newline at end of file diff --git a/scripts/main_exp/gen_data.sh b/scripts/main_exp/gen_data.sh index c1c0a26..78e7eb2 100755 --- a/scripts/main_exp/gen_data.sh +++ b/scripts/main_exp/gen_data.sh @@ -16,5 +16,5 @@ done uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --glob_pattern 'data/raw_datasets/fw_qa_v2/min_0_to_2000/*_level_0_val.parquet' # self-gen data for other ds -uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --ds_names squad_compact ropes_compact drop_compact --closed_qa_prob 1.0 -uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --ds_names pwc_compact --closed_qa_prob 0.0 +uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --ds_names squad_compact ropes_compact drop_compact --split train --closed_qa_prob 1.0 +uv run data/self_generate_qa.py --vllm_model google/gemma-2-2b-it --ds_names pwc_compact --split train --closed_qa_prob 0.0 diff --git a/webui/SELF_GEN_VIEWER.md b/webui/SELF_GEN_VIEWER.md new file mode 100644 index 0000000..408b40d --- /dev/null +++ b/webui/SELF_GEN_VIEWER.md @@ -0,0 +1,31 @@ +# Self-Gen Data Viewer + +Thanks Claude. + +Running the viewer +```bash +uv run self_gen_viewer.py +``` + +Then open your browser and go to: **http://localhost:5001** + +## Usage + +1. **Select a Model Folder**: Choose from the dropdown list (e.g., `google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0`) +2. **Select a Parquet File**: Once a folder is selected, available parquet files will appear +3. **Set Number of Samples**: Adjust the sample count (default: 100, max: 1000) +4. **Click "Load Data"**: View the visualized data with context and Q&A pairs + +## Data Structure + +The viewer expects data in the following structure: +``` +data/raw_datasets/self_gen/ +├── google/ +│ └── gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0/ +│ └── fw_qa_v2/ +│ └── *.parquet +└── mistralai/ + └── Mistral-7B-Instruct-v0.2_temp_0.0_closed_qa_prob_1.0/ + └── *.parquet +``` diff --git a/webui/self_gen_viewer.py b/webui/self_gen_viewer.py new file mode 100644 index 0000000..6c8cd13 --- /dev/null +++ b/webui/self_gen_viewer.py @@ -0,0 +1,170 @@ +import traceback +from pathlib import Path + +from datasets import load_dataset +from flask import Flask, jsonify, render_template, request +from transformers import AutoTokenizer + +app = Flask(__name__) + +# Base path for self_gen data +BASE_DATA_PATH = Path(__file__).parent.parent / "data" / "raw_datasets" / "self_gen" + +# Cache for tokenizers +tokenizer_cache = {} + + +def get_tokenizer(model_path): + """Get or create tokenizer with caching""" + if model_path not in tokenizer_cache: + try: + tokenizer_cache[model_path] = AutoTokenizer.from_pretrained(model_path) + except Exception as e: + print(f"Error loading tokenizer for {model_path}: {e}") + return None + return tokenizer_cache[model_path] + + +def discover_folders(): + """Discover all model folders in self_gen directory""" + folders = [] + if not BASE_DATA_PATH.exists(): + return folders + + for vendor_dir in BASE_DATA_PATH.iterdir(): + if vendor_dir.is_dir(): + for model_dir in vendor_dir.iterdir(): + if model_dir.is_dir(): + rel_path = model_dir.relative_to(BASE_DATA_PATH) + folders.append(str(rel_path)) + + return sorted(folders) + + +def discover_parquet_files(folder_path): + """Discover all parquet files in a folder""" + full_path = BASE_DATA_PATH / folder_path + parquet_files = [] + + if full_path.exists(): + for parquet_file in full_path.glob("**/*.parquet"): + rel_path = parquet_file.relative_to(full_path) + parquet_files.append(str(rel_path)) + + return sorted(parquet_files) + + +def extract_model_name_from_folder(folder_path): + """Extract base model name from folder path""" + # e.g., "google/gemma-2-2b-it_temp_0.0_closed_qa_prob_1.0" -> "google/gemma-2-2b-it" + parts = folder_path.split("/") + if len(parts) >= 2: + vendor = parts[0] + model_part = parts[1].split("_temp_")[0] + return f"{vendor}/{model_part}" + return None + + +@app.route("/") +def index(): + """Main page""" + folders = discover_folders() + return render_template("self_gen_viewer.html", folders=folders) + + +@app.route("/api/folders") +def api_folders(): + """API endpoint to get available folders""" + folders = discover_folders() + return jsonify({"folders": folders}) + + +@app.route("/api/parquet_files") +def api_parquet_files(): + """API endpoint to get parquet files in a folder""" + folder = request.args.get("folder", "") + if not folder: + return jsonify({"error": "No folder specified"}), 400 + + files = discover_parquet_files(folder) + return jsonify({"files": files}) + + +@app.route("/api/load_data") +def api_load_data(): + """API endpoint to load and display data from a parquet file""" + folder = request.args.get("folder", "") + parquet_file = request.args.get("file", "") + num_samples = int(request.args.get("num_samples", 100)) + + if not folder or not parquet_file: + return jsonify({"error": "Missing parameters"}), 400 + + try: + # Construct full path + full_path = BASE_DATA_PATH / folder / parquet_file + + if not full_path.exists(): + return jsonify({"error": f"File not found: {full_path}"}), 404 + + # Extract model name for tokenizer + model_name = extract_model_name_from_folder(folder) + if not model_name: + return jsonify({"error": "Could not extract model name from folder"}), 400 + + # Load tokenizer + tokenizer = get_tokenizer(model_name) + if tokenizer is None: + return jsonify({"error": f"Could not load tokenizer for {model_name}"}), 500 + + # Load dataset + ds = load_dataset( + "parquet", data_files=str(full_path), split=f"train[:{num_samples}]" + ) + + # Process samples + samples = [] + for i, sample in enumerate(ds): + processed_sample = { + "index": i, + "ctx": tokenizer.decode(sample["ctx_ids"], skip_special_tokens=False) + if "ctx_ids" in sample + else "N/A", + "questions": [], + } + + # Decode input_ids if present + if "input_ids" in sample: + if isinstance(sample["input_ids"][0], list): + # Multiple Q&A pairs + processed_sample["questions"] = [ + tokenizer.decode(qa, skip_special_tokens=False) + for qa in sample["input_ids"] + ] + else: + # Single item + processed_sample["questions"] = [ + tokenizer.decode(sample["input_ids"], skip_special_tokens=False) + ] + + samples.append(processed_sample) + + return jsonify( + { + "success": True, + "num_samples": len(samples), + "model_name": model_name, + "file_path": str(parquet_file), + "samples": samples, + } + ) + + except Exception as e: + traceback.print_exc() + return jsonify({"error": str(e), "traceback": traceback.format_exc()}), 500 + + +if __name__ == "__main__": + print(f"Data path: {BASE_DATA_PATH}") + print(f"Available folders: {discover_folders()}") + app.run(debug=True, host="0.0.0.0", port=5001) diff --git a/webui/templates/self_gen_viewer.html b/webui/templates/self_gen_viewer.html new file mode 100644 index 0000000..562d31d --- /dev/null +++ b/webui/templates/self_gen_viewer.html @@ -0,0 +1,422 @@ + + + + + + + Self-Gen Data Viewer + + + + +
+

🔍 Self-Gen Data Viewer

+ +
+
+ + +
+ +
+ + +
+ +
+
+ + +
+ +
+
+ + +
+
+ + + + + \ No newline at end of file