|
import pandas as pd |
|
import matplotlib.pyplot as plt |
|
import seaborn as sns |
|
import json |
|
import argparse |
|
import os |
|
from scipy.stats import gaussian_kde |
|
import numpy as np |
|
|
|
def get_model_for(doc_type: str, override_model: str) -> str: |
|
"""Returns model type or the override model if specified""" |
|
if override_model: |
|
return override_model |
|
doc_type = doc_type.split("_", 1)[0] |
|
if doc_type in ("book", "books", "pg19"): |
|
return "books_pp" |
|
elif doc_type in ("culturax", "slimpajama", "wikipedia", "digimanus"): |
|
return "wikipedia_pp" |
|
elif doc_type in ("newspaper", "newspapers"): |
|
return "newspapers_pp" |
|
elif doc_type in ("evalueringsrapport", "lovdata", "maalfrid", "parlamint"): |
|
return "maalfrid_pp" |
|
else: |
|
return "wikipedia_pp" |
|
|
|
def load_data(files): |
|
all_data = [] |
|
for file_path in files: |
|
with open(file_path, 'r') as file: |
|
lines = file.readlines() |
|
data = [json.loads(line) for line in lines] |
|
all_data.extend(data) |
|
return pd.DataFrame(all_data) |
|
|
|
def plot_histograms(files, output_folder, xlim, override_model): |
|
df = load_data(files) |
|
doc_types = df['doctype'].unique() |
|
fig, axes = plt.subplots(len(doc_types), 1, figsize=(12, 4 * len(doc_types)), squeeze=False) |
|
|
|
|
|
palette = sns.color_palette("husl", len(doc_types)) |
|
|
|
for i, doc_type in enumerate(doc_types): |
|
ax = axes[i, 0] |
|
group = df[df['doctype'] == doc_type] |
|
languages = group['lang'].unique() |
|
|
|
|
|
colors = sns.color_palette("husl", len(languages)) |
|
|
|
for j, lang in enumerate(languages): |
|
lang_group = group[group['lang'] == lang] |
|
perplexity_model = get_model_for(doc_type, override_model) |
|
perplexity_values = lang_group['perplexities'].apply(lambda x: x[perplexity_model]).values |
|
|
|
series_color = colors[j] |
|
|
|
|
|
sns.histplot(perplexity_values, ax=ax, color=series_color, alpha=0.3, element="step", fill=True, stat="density", binwidth=30) |
|
|
|
|
|
sns.kdeplot(perplexity_values, ax=ax, bw_adjust=2, color=series_color, label=f"{lang} - {doc_type} ({perplexity_model})", linewidth=1.5) |
|
|
|
|
|
kde = gaussian_kde(perplexity_values) |
|
x_range = np.linspace(0, xlim, 1000) |
|
y_values = kde.evaluate(x_range) |
|
|
|
quartiles = np.quantile(perplexity_values, [0.25, 0.5, 0.75]) |
|
quartile_labels = ["Q1", "Q2", "Q3"] |
|
for q, quartile in enumerate(quartiles): |
|
idx = (np.abs(x_range-quartile)).argmin() |
|
y_quartile = y_values[idx] |
|
ax.plot([quartile, quartile], [0, y_quartile], color=series_color, linestyle='--', linewidth=1) |
|
ax.text(quartile, y_quartile, f'{quartile_labels[q]}: {quartile:.2f}', verticalalignment='bottom', horizontalalignment='right', color=series_color, fontsize=6) |
|
|
|
ax.set_title(f'Document Type: {doc_type} ({perplexity_model})') |
|
ax.set_xlabel('Perplexity Value') |
|
ax.set_ylabel('Density') |
|
ax.legend() |
|
ax.set_xlim(left=0, right=xlim) |
|
|
|
plt.tight_layout() |
|
output_filename = os.path.join(output_folder, "all_doc_types_plots.png") |
|
plt.savefig(output_filename, dpi=300) |
|
plt.close(fig) |
|
print(f"All document type plots saved to {output_filename}") |
|
|
|
def main(): |
|
parser = argparse.ArgumentParser(description="Plot histograms from JSON lines files.") |
|
parser.add_argument('files', nargs='+', help="Path to the JSON lines files") |
|
parser.add_argument('-o', '--output_folder', default=".", help="Output folder for the plots") |
|
parser.add_argument('--xlim', type=int, default=2500, help="Maximum x-axis limit for the plots") |
|
parser.add_argument('--model', default="", help="Override the perplexity model for all plots") |
|
|
|
args = parser.parse_args() |
|
|
|
if not os.path.exists(args.output_folder): |
|
os.makedirs(args.output_folder, exist_ok=True) |
|
|
|
plot_histograms(args.files, args.output_folder, args.xlim, args.model) |
|
|
|
if __name__ == "__main__": |
|
main() |
|
|
