qingy2024

Training in progress, step 8900

2204d6b verified about 1 year ago

15.7 kB

	import os
	import re
	import numpy as np
	import time
	import shutil
	import json
	import matplotlib.pyplot as plt
	from huggingface_hub import login, create_repo, upload_folder, HfFolder
	from pathlib import Path # Using pathlib for easier path manipulation

	# --- Configuration Constants ---
	# Model and Repo Details
	BASE_MODEL_NAME = "Qwen/Qwen2.5-3B-Instruct"
	TARGET_REPO_NAME = "Tesslate/Gradience-T1-3B-Checkpoint" # Specify your target repo

	# Training Parameters (Update if necessary)
	TOTAL_STEPS = 9838 # Total expected steps for progress calculation

	# File Names
	README_FILENAME = "README.md"
	ADAPTER_CONFIG_FILENAME = "adapter_config.json"
	TRAINER_STATE_FILENAME = "trainer_state.json"
	LOSS_PLOT_FILENAME = "loss.png"

	# Plotting Configuration
	LOSS_SMOOTHING_WINDOW = 40

	# Monitoring Configuration
	CHECKPOINT_DIR_PATTERN = re.compile(r"^checkpoint-(\d+)$")
	POLL_INTERVAL_SECONDS = 30
	PRE_UPLOAD_DELAY_SECONDS = 10 # Delay after finding checkpoint before processing

	# --- Global State ---
	# Set to track uploaded checkpoints (using Path objects for consistency)
	uploaded_checkpoints = set()

	# --- Helper Functions ---

	def get_huggingface_token():
	"""Retrieves the Hugging Face token from environment variable or login cache."""
	token = os.getenv('HUGGINGFACE_TOKEN')
	if token:
	print("Using Hugging Face token from HUGGINGFACE_TOKEN environment variable.")
	return token
	token = HfFolder.get_token()
	if token:
	print("Using Hugging Face token from saved credentials.")
	return token
	raise ValueError("Hugging Face token not found. Set HUGGINGFACE_TOKEN environment variable or login using `huggingface-cli login`.")

	def update_adapter_config(config_path: Path, base_model_name: str):
	"""
	Reads adapter_config.json, updates the base_model_name_or_path field,
	and saves it back.

	Args:
	config_path (Path): Path to the adapter_config.json file.
	base_model_name (str): The base model name to set.
	"""
	try:
	with open(config_path, 'r') as file:
	config = json.load(file)

	config['base_model_name_or_path'] = base_model_name

	with open(config_path, 'w') as file:
	json.dump(config, file, indent=2)
	print(f"Updated 'base_model_name_or_path' in {config_path}")

	except FileNotFoundError:
	print(f"Error: Adapter config file not found at {config_path}")
	except json.JSONDecodeError:
	print(f"Error: Could not decode JSON from {config_path}. Is it valid?")
	except KeyError:
	print(f"Error: 'base_model_name_or_path' key not found in {config_path}")
	except Exception as e:
	print(f"An unexpected error occurred while updating {config_path}: {e}")

	def generate_readme_content(checkpoint_number: int, total_steps: int, base_model: str, loss_plot_filename: str) -> str:
	"""Generates the README content with updated progress."""
	if total_steps <= 0:
	progress_percentage = 0.0
	else:
	progress_percentage = min(100.0, (checkpoint_number / total_steps) * 100) # Ensure percentage doesn't exceed 100

	progress_width = f"{progress_percentage:.2f}%"
	progress_text = f"Progress: {checkpoint_number} out of {total_steps} steps"

	# Using an f-string for the template makes insertions cleaner
	readme_template = f"""
	---
	base_model: {base_model}
	library_name: peft
	---
	# Gradience T1 3B (Step {checkpoint_number} Checkpoint)

	> [!NOTE]
	> Training in progress...

	<!DOCTYPE html>
	<html lang="en">
	<head>
	<meta charset="UTF-8">
	<title>Progress Bar Example</title>
	<style>
	.progress-container {{
	width: 100%;
	background-color: #e0e0e0;
	border-radius: 25px;
	overflow: hidden;
	margin: 20px 0;
	}}
	.progress-bar {{
	height: 30px;
	width: 0;
	background-color: #76c7c0;
	text-align: center;
	line-height: 30px;
	color: white;
	border-radius: 25px 0 0 25px;
	}}
	.progress-text {{
	margin-top: 10px;
	font-size: 16px;
	font-family: Arial, sans-serif;
	}}
	</style>
	</head>
	<body>
	<div style="width: 100%; background-color: #e0e0e0; border-radius: 25px; overflow: hidden; margin: 20px 0;">
	<div style="height: 30px; width: {progress_width}; background-color: #76c7c0; text-align: center; line-height: 30px; color: white; border-radius: 25px 0 0 25px;">
	<!-- {progress_percentage:.2f}% -->
	</div>
	</div>
	<p style="font-family: Arial, sans-serif; font-size: 16px;">{progress_text}</p>
	</body>
	</html>

	## Training Loss
	![Training Loss Chart]({loss_plot_filename})
	""".strip()
	return readme_template

	def plot_loss_from_json(
	json_file_path: Path,
	output_image_path: Path,
	smooth_steps: int = LOSS_SMOOTHING_WINDOW
	):
	"""
	Reads training log data from a JSON file (trainer_state.json),
	extracts loss and step values, plots the original loss and a smoothed
	version (running average), and saves the plot to a PNG file.

	Args:
	json_file_path (Path): Path to the input trainer_state.json file.
	output_image_path (Path): Path where the output PNG plot will be saved.
	smooth_steps (int): Window size for running average smoothing.
	If <= 0, no smoothing is applied.
	"""
	print(f"Reading training log data from: {json_file_path}")
	print(f"Smoothing window: {smooth_steps if smooth_steps > 0 else 'Disabled'}")

	try:
	with open(json_file_path, 'r') as f:
	data = json.load(f)
	except FileNotFoundError:
	print(f"Error: JSON file not found at {json_file_path}")
	return
	except json.JSONDecodeError:
	print(f"Error: Could not decode JSON from {json_file_path}. Is it valid?")
	return
	except Exception as e:
	print(f"An unexpected error occurred while reading {json_file_path}: {e}")
	return

	log_history = data.get("log_history") # Use .get for safer access
	if not isinstance(log_history, list):
	print(f"Error: 'log_history' key not found or not a list in {json_file_path}")
	return

	steps, losses = [], []
	for entry in log_history:
	if isinstance(entry, dict) and "step" in entry and "loss" in entry and entry["loss"] is not None:
	try:
	steps.append(int(entry["step"]))
	losses.append(float(entry["loss"]))
	except (ValueError, TypeError):
	print(f"Warning: Skipping entry with non-numeric step/loss: {entry}")
	# else: # Optionally log skipped entries
	# print(f"Info: Skipping log entry missing 'step'/'loss' or loss is null: {entry}")

	if not steps:
	print("No valid step/loss data found in the log history to plot.")
	return

	# Convert to numpy arrays and sort by step (good practice)
	steps = np.array(steps)
	losses = np.array(losses)
	sorted_indices = np.argsort(steps)
	steps = steps[sorted_indices]
	losses = losses[sorted_indices]

	print(f"Found {len(steps)} valid data points to plot.")

	# Calculate Running Average
	smoothed_losses = None
	smoothed_steps = None
	apply_smoothing = smooth_steps > 0 and len(losses) >= smooth_steps

	if apply_smoothing:
	try:
	weights = np.ones(smooth_steps) / smooth_steps
	smoothed_losses = np.convolve(losses, weights, mode='valid')
	smoothed_steps = steps[smooth_steps - 1:] # Steps corresponding to the smoothed values
	print(f"Calculated smoothed loss over {len(smoothed_steps)} points.")
	except Exception as e:
	print(f"Warning: Could not calculate smoothed loss. Error: {e}")
	apply_smoothing = False # Disable if calculation fails
	elif smooth_steps > 0:
	print(f"Warning: Not enough data points ({len(losses)}) for smoothing window ({smooth_steps}). Skipping smoothing.")

	# Plotting
	plt.style.use('seaborn-v0_8-darkgrid') # Use a nice style
	plt.figure(figsize=(10, 6)) # Standard figure size

	plt.plot(steps, losses, linestyle='-', color='skyblue', alpha=0.5, label='Original Loss')

	if apply_smoothing and smoothed_losses is not None and smoothed_steps is not None:
	plt.plot(smoothed_steps, smoothed_losses, linestyle='-', color='dodgerblue', alpha=1.0, linewidth=1.5,
	label=f'Smoothed Loss ({smooth_steps}-step avg)')

	plt.xlabel("Step")
	plt.ylabel("Loss")
	plt.title("Training Loss Progression")
	plt.legend()
	plt.tight_layout() # Adjust layout

	# Saving
	try:
	plt.savefig(output_image_path, format='png', dpi=150)
	print(f"Plot successfully saved to: {output_image_path}")
	except Exception as e:
	print(f"Error saving plot to {output_image_path}: {e}")
	finally:
	plt.close() # Ensure figure is closed to free memory

	def prepare_checkpoint_folder(checkpoint_path: Path, checkpoint_number: int):
	"""
	Updates README.md, adapter_config.json, and generates the loss plot
	within the specified checkpoint folder.
	"""
	print(f"Preparing checkpoint folder: {checkpoint_path}")

	# 1. Update adapter config
	adapter_config_path = checkpoint_path / ADAPTER_CONFIG_FILENAME
	update_adapter_config(adapter_config_path, BASE_MODEL_NAME)

	# 2. Generate loss plot
	trainer_state_path = checkpoint_path / TRAINER_STATE_FILENAME
	loss_plot_path = checkpoint_path / LOSS_PLOT_FILENAME
	plot_loss_from_json(trainer_state_path, loss_plot_path, smooth_steps=LOSS_SMOOTHING_WINDOW)

	# 3. Generate and write README
	readme_path = checkpoint_path / README_FILENAME
	readme_content = generate_readme_content(checkpoint_number, TOTAL_STEPS, BASE_MODEL_NAME, LOSS_PLOT_FILENAME)
	try:
	with open(readme_path, 'w', encoding='utf-8') as file:
	file.write(readme_content)
	print(f"Generated and saved {README_FILENAME} in {checkpoint_path}")
	except Exception as e:
	print(f"Error writing README file to {readme_path}: {e}")

	# --- Core Logic ---

	def find_new_checkpoint(current_dir: Path = Path('.')) -> tuple[int, Path] \| None:
	"""
	Finds the checkpoint folder in the specified directory with the highest
	step number that has not been previously uploaded.

	Args:
	current_dir (Path): The directory to scan for checkpoints.

	Returns:
	tuple[int, Path] \| None: A tuple containing the (checkpoint_number, folder_path)
	or None if no new checkpoint is found.
	"""
	new_checkpoints = []
	try:
	for item in current_dir.iterdir():
	if item.is_dir():
	match = CHECKPOINT_DIR_PATTERN.match(item.name)
	# Check if it matches the pattern AND has not been uploaded
	if match and item not in uploaded_checkpoints:
	checkpoint_number = int(match.group(1))
	new_checkpoints.append((checkpoint_number, item))
	except FileNotFoundError:
	print(f"Error: Directory not found: {current_dir}")
	return None
	except Exception as e:
	print(f"Error scanning directory {current_dir}: {e}")
	return None

	if new_checkpoints:
	new_checkpoints.sort(key=lambda x: x[0], reverse=True) # Sort by step number, highest first
	return new_checkpoints[0] # Return the one with the highest step number
	return None

	def upload_checkpoint_to_hf(folder_path: Path, checkpoint_number: int, repo_id: str):
	"""
	Uploads the prepared checkpoint folder to Hugging Face Hub and deletes
	the folder locally upon successful upload.

	Args:
	folder_path (Path): Path to the local checkpoint folder.
	checkpoint_number (int): The checkpoint step number.
	repo_id (str): The Hugging Face repository ID (e.g., "username/repo-name").
	"""
	print(f"\nAttempting to upload {folder_path.name} to Hugging Face repository: {repo_id}...")

	try:
	# Ensure repository exists
	create_repo(repo_id, repo_type="model", exist_ok=True)
	print(f"Repository {repo_id} exists or was created.")

	# Upload the folder contents
	upload_folder(
	folder_path=str(folder_path), # upload_folder expects string path
	repo_id=repo_id,
	commit_message=f"Upload checkpoint {checkpoint_number}",
	repo_type="model" # Explicitly set repo type
	)
	print(f"Successfully uploaded contents of {folder_path.name} to {repo_id}.")

	# Delete the local folder ONLY after successful upload
	try:
	shutil.rmtree(folder_path)
	print(f"Successfully deleted local folder: {folder_path}")
	return True # Indicate success
	except OSError as e:
	print(f"Error deleting local folder {folder_path}: {e}. Please delete manually.")
	return True # Upload succeeded, but deletion failed

	except Exception as e:
	print(f"ERROR during Hugging Face upload for {folder_path.name}: {e}")
	print("Upload failed. Local folder will not be deleted.")
	return False # Indicate failure

	# --- Main Execution ---

	def main():
	"""
	Main loop to monitor for new checkpoints, prepare them, upload them to
	Hugging Face Hub, and clean up locally.
	"""
	try:
	hf_token = get_huggingface_token()
	login(hf_token)
	print("\nSuccessfully logged into Hugging Face Hub.")
	except ValueError as e:
	print(f"Error: {e}")
	return # Exit if login fails
	except Exception as e:
	print(f"An unexpected error occurred during Hugging Face login: {e}")
	return

	print("\nStarting checkpoint monitor...")
	print(f"Will check for new checkpoints matching '{CHECKPOINT_DIR_PATTERN.pattern}' every {POLL_INTERVAL_SECONDS} seconds.")
	print(f"Target repository: {TARGET_REPO_NAME}")
	print(f"Found checkpoints will be tracked (not re-uploaded): {uploaded_checkpoints or 'None yet'}")
	print("-" * 30)

	while True:
	new_checkpoint_info = find_new_checkpoint()

	if new_checkpoint_info:
	checkpoint_number, folder_path = new_checkpoint_info
	print(f"\nFound new checkpoint: {folder_path.name} (Step {checkpoint_number})")

	# Optional delay: wait a bit in case files are still being written
	print(f"Waiting {PRE_UPLOAD_DELAY_SECONDS} seconds before processing...")
	time.sleep(PRE_UPLOAD_DELAY_SECONDS)

	# Prepare the folder (update README, config, generate plot)
	prepare_checkpoint_folder(folder_path, checkpoint_number)

	# Attempt upload and deletion
	upload_successful = upload_checkpoint_to_hf(
	folder_path=folder_path,
	checkpoint_number=checkpoint_number,
	repo_id=TARGET_REPO_NAME
	)

	if upload_successful:
	# Add to uploaded set ONLY if upload (and optionally deletion) was processed
	uploaded_checkpoints.add(folder_path)
	print(f"Added {folder_path.name} to the set of processed checkpoints.")

	print("-" * 30) # Separator after processing a checkpoint

	else:
	# Use \r for inline update when no checkpoint found
	print(f"\rNo new checkpoints found. Checking again in {POLL_INTERVAL_SECONDS} seconds... ", end="")

	# Wait before the next check
	time.sleep(POLL_INTERVAL_SECONDS)

	if __name__ == "__main__":
	try:
	main()
	except KeyboardInterrupt:
	print("\nMonitoring stopped by user.")