update(core): fix code

app.py

@@ -1,69 +1,83 @@
 import gradio as gr
 
-from utils import checkpoints, load_model, log_perplexity
+from utils import load_model, log_perplexity
 
 
-class ModelManager:
+class Manager:
     """Class to manage model loading and perplexity calculation state."""
 
     def __init__(self):
-        self.loaded_models = None
+        self.loaded = None
 
     def load_models(self, checkpoint_input_str: str) -> str:
         """Load models from a comma-separated string of checkpoint names."""
-        checkpoint_list = [
-            c.strip() for c in checkpoint_input_str.split(",") if c.strip()
+        checkpoints = [
+            ckpt.strip() for ckpt in checkpoint_input_str.split(",") if ckpt.strip()
         ]
 
-        if not checkpoint_list:
+        if not checkpoints:
             return "Please enter at least one model checkpoint name."
 
         try:
-            self.loaded_models = load_model(checkpoint_list)
+            self.loaded = load_model(checkpoints)
             return "Models loaded successfully!"
         except Exception as e:
             return f"Model loading failed: {e}"
 
-    def calculate_perplexity(self) -> dict | str:
+    def perplexity(
+        self,
+        num_samples: int | None = None,
+        sample_length: int | None = None,
+    ) -> dict | str:
         """Calculate perplexity using the loaded models."""
-        if self.loaded_models is None:
+        if self.loaded is None:
             return "Please load models first."
-
+        if num_samples is None or sample_length is None:
+            return "Please set the number of samples and sample length."
         try:
-            result = log_perplexity()
-            return result
+            return log_perplexity(self.loaded, num_samples, sample_length)
         except Exception as e:
             return f"Perplexity calculation failed: {e}"
 
 
-def create_interface() -> gr.Blocks:
+def make_interface() -> gr.Blocks:
     """Create and return the Gradio interface."""
-    manager = ModelManager()
+    manager = Manager()
 
     with gr.Blocks() as demo:
-        gr.Markdown("# LLM PPL")
+        gr.Markdown("# LLM PPLs")
 
-        checkpoint_input = gr.Textbox(
-            label="Checkpoints",
-            value=", ".join(checkpoints),
+        checkpoints = gr.Textbox(
+            label="Checkpoints", value="HuggingFaceTB/SmolLM2-135M"
         )
 
         load_btn = gr.Button("Load Models", variant="primary")
-        perplexity_btn = gr.Button("Compute PPL")
 
-        load_output = gr.Textbox(label="Model Loading Status", interactive=False)
+        with gr.Row():
+            num_samples = gr.Number(label="Number of Samples", value=1500)
+            sample_length = gr.Number(label="Sample Length", value=128)
+
+        perplexity_btn = gr.Button("Compute PPLs")
+
+        load_output = gr.Textbox(label="Model Loading Status")
         perplexity_output = gr.JSON(label="PPL Results")
 
         # Connect event handlers
         load_btn.click(
-            fn=manager.load_models, inputs=checkpoint_input, outputs=load_output
+            fn=manager.load_models,
+            inputs=checkpoints,
+            outputs=load_output,
         )
 
-        perplexity_btn.click(fn=manager.calculate_perplexity, outputs=perplexity_output)
+        perplexity_btn.click(
+            fn=manager.perplexity,
+            inputs=[num_samples, sample_length],
+            outputs=perplexity_output,
+        )
 
     return demo
 
 
 if __name__ == "__main__":
-    demo = create_interface()
+    demo = make_interface()
     demo.launch()

utils.py

@@ -1,94 +1,78 @@
 from typing import Final
 
+import gradio as gr
 import numpy as np
 import torch
 import ujson as json
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-dev: Final = "cuda" if torch.cuda.is_available() else "cpu"
-texts: Final = json.load(open("texts.json", "r"))
+_dev: Final = "cuda" if torch.cuda.is_available() else "cpu"
+_dtype: Final = torch.bfloat16
 
-checkpoints = ["HuggingFaceTB/SmolLM2-135M"]  # Inputs
 
+def _perplexity(model, tokenizer, text) -> float:
+    encodings = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
+    input_ids = encodings.input_ids.to(_dev)
+    with torch.no_grad():
+        outputs = model(input_ids, labels=input_ids)
+        loss = outputs.loss.item()
+    return np.log(torch.exp(torch.tensor(loss)).item())
 
-def load_model(checkpoints: list[str]) -> dict:
-    tokenizers = [
-        AutoTokenizer.from_pretrained(checkpoint) for checkpoint in checkpoints
-    ]
 
+def load_model(checkpoints: list[str]) -> dict:
+    tokenizers = [AutoTokenizer.from_pretrained(c) for c in checkpoints]
     models = [
-        AutoModelForCausalLM.from_pretrained(
-            checkpoint,
-            device_map="auto",
-            torch_dtype=torch.bfloat16,
-        )
-        .to(dev)
-        .eval()
-        for checkpoint in checkpoints
+        AutoModelForCausalLM.from_pretrained(c, device_map="auto", torch_dtype=_dtype)
+        for c in checkpoints
     ]
 
     # Load the models and tokenizers into a dictionary
     return {
-        checkpoint: {"model": model, "tokenizer": tokenizer}
-        for checkpoint, model, tokenizer in zip(checkpoints, models, tokenizers)
+        ckpt: {"model": model.to(_dev).eval(), "tokenizer": tokenizer}
+        for ckpt, model, tokenizer in zip(checkpoints, models, tokenizers)
     }
 
 
-def _perplexity(model, tokenizer, text):
-    encodings = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
-    input_ids = encodings.input_ids.to(dev)
-    with torch.no_grad():
-        outputs = model(input_ids, labels=input_ids)
-        loss = outputs.loss.item()
-    return torch.exp(torch.tensor(loss)).item()
-
-
-num_samples: Final[int] = 500  # Sample size for perplexity calculation
-sample_length: Final[int] = 100  # Maximum length of text to consider for perplexity
-
-loaded = load_model(checkpoints)
+def log_perplexity(
+    loaded: dict,
+    num_samples: int,
+    sample_length: int,
+    progress=gr.Progress(),
+) -> dict:
+    # Initialize a dictionary to store perplexity
+    ppls: dict[str, list] = {ckpt: [] for ckpt in loaded.keys()}
 
+    # Initialize samples
+    texts: Final[list[str]] = [
+        text.strip()[:sample_length]
+        for text in json.load(open("texts.json", "r"))
+        if text.strip()
+    ]
 
-def log_perplexity() -> dict:
-    # Initialize a dictionary to store perplexity
-    ppls = {checkpoint: [] for checkpoint in loaded.keys()}
+    # Start the iteration
+    progress(0, desc="Starting")
     for i in range(num_samples):
-        text = texts[i]
-        if len(text.strip()) == 0:
-            continue
-
-        text = text.strip()[:sample_length]
-
-        # Calculate perplexity for each model
-        current_ppls = {}
-        for checkpoint, info in loaded.items():
-            ppl = _perplexity(
-                info["model"],
-                info["tokenizer"],
-                text,
-            )
-            current_ppls[checkpoint] = ppl
-
-        # Filter out outliers
-        if all(1 < ppl < 1e4 for ppl in current_ppls.values()):
-            for checkpoint, ppl in current_ppls.items():
-                ppls[checkpoint].append(ppl)
-
-    # Convert perplexity into log scale
-    log_ppls: dict = {checkpoint: np.log(ppl) for checkpoint, ppl in ppls.items()}
+        progress(i / num_samples, desc="Processing samples")
+        for ckpt, info in loaded.items():  # Calculate perplexity for each model
+            ppl: float = _perplexity(info["model"], info["tokenizer"], texts[i])
+            if 1 < ppl < 1e4:  # Filter out outliers
+                ppls[ckpt].append(ppl)
 
     # Calculate the mean perplexity for each model
-    mean_log_ppls: dict = {
-        checkpoint: np.mean(ppl) for checkpoint, ppl in log_ppls.items()
-    }
+    means: dict = {ckpt: np.mean(ppl) for ckpt, ppl in ppls.items()}
 
     # Calculate the standard deviation of perplexity for each model
-    std_log_ppls: dict = {
-        checkpoint: np.std(ppl) for checkpoint, ppl in log_ppls.items()
-    }
+    stds: dict = {ckpt: np.std(ppl) for ckpt, ppl in ppls.items()}
 
-    return {
-        "ppls": ppls,
-        "mean_ppls": mean_log_ppls,
-        "std_ppls": std_log_ppls,
-    }
+    return {"ppls": ppls, "means": means, "stds": stds}
+
+
+if __name__ == "__main__":
+    from pprint import pprint
+
+    # Example usage
+    checkpoints = ["HuggingFaceTB/SmolLM2-135M"]
+    loaded = load_model(checkpoints)
+    num_samples = 500
+    sample_length = 128
+    pprint(log_perplexity(loaded, num_samples, sample_length))