import json

# Path to the dataset_info.json file
file_path = "LLaMA-Factory/data/dataset_info.json"

# Load the existing JSON data
with open(file_path, 'r') as file:
    dataset_info = json.load(file)

# Create the new entry for your dataset
new_entry = {
    "forecasting_sft": {
        "hf_hub_url": "shash42/deepseek-r1-forecasting-sft",
        "formatting": "sharegpt",
        "columns": {
            "messages": "messages"
        },
        "tags": {
            "role_tag": "role",
            "content_tag": "content",
            "user_tag": "user",
            "assistant_tag": "assistant",
            "system_tag": "system"
        }
    }
}

# Add the new entry to the dataset_info
dataset_info.update(new_entry)

# Write the updated data back to the file
with open(file_path, 'w') as file:
    json.dump(dataset_info, file, indent=2)

print(f"Successfully added 'forecasting_sft' to {file_path}")