import json
import matplotlib.pyplot as plt

# Step 1: Load the data from the JSONL file
file_path = '/root/autodl-tmp/fhy/finetune/mistral-nosysprompt/output/mistral-lora-fp16-epoch4/trainer_log.jsonl'  # Update this with the correct file path

# Read the content of the JSONL file
with open(file_path, 'r') as file:
    json_lines = file.readlines()

# Parse each line as a JSON object
data = [json.loads(line) for line in json_lines]

# Step 2: Extract steps and loss values
steps = [entry['current_steps'] for entry in data]
loss = [entry['loss'] for entry in data]

# Step 3: Plot the data as a simple line plot without markers
plt.figure(figsize=(10, 6))
plt.plot(steps, loss, linestyle='-', color='b')

# Add labels and title
plt.xlabel('Steps')
plt.ylabel('Loss')
plt.title('Loss vs Steps')

# Display the plot
plt.grid(True)
plt.savefig('/root/autodl-tmp/fhy/finetune/mistral-nosysprompt/output/mistral-lora-fp16-epoch4/loss_vs_steps.png')