chore(misc): update files

app.py

@@ -1,7 +1,70 @@
 import gradio as gr
 
-def greet(name):
-    return "Hello " + name + "!!"
+from utils import checkpoints, load_model, log_perplexity
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+
+class ModelManager:
+    """Class to manage model loading and perplexity calculation state."""
+
+    def __init__(self):
+        self.loaded_models = 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()
+        ]
+
+        if not checkpoint_list:
+            return "Please enter at least one model checkpoint name."
+
+        try:
+            self.loaded_models = load_model(checkpoint_list)
+            return "Models loaded successfully!"
+        except Exception as e:
+            return f"Model loading failed: {e}"
+
+    def calculate_perplexity(self) -> dict | str:
+        """Calculate perplexity using the loaded models."""
+        if self.loaded_models is None:
+            return "Please load models first."
+
+        try:
+            result = log_perplexity()
+            return result
+        except Exception as e:
+            return f"Perplexity calculation failed: {e}"
+
+
+def create_interface() -> gr.Blocks:
+    """Create and return the Gradio interface."""
+    manager = ModelManager()
+
+    with gr.Blocks() as demo:
+        gr.Markdown("# Language Model Perplexity Calculator (Sequential Version)")
+        gr.Markdown("Enter Hugging Face model checkpoint names (comma-separated).")
+
+        checkpoint_input = gr.Textbox(
+            label="Checkpoints (e.g. Qwen/Qwen2.5-14B-Instruct)",
+            value=", ".join(checkpoints),
+        )
+
+        load_btn = gr.Button("Load Models", variant="primary")
+        perplexity_btn = gr.Button("Calculate Perplexity")
+
+        load_output = gr.Textbox(label="Model Loading Status", interactive=False)
+        perplexity_output = gr.JSON(label="Perplexity Results")
+
+        # Connect event handlers
+        load_btn.click(
+            fn=manager.load_models, inputs=checkpoint_input, outputs=load_output
+        )
+
+        perplexity_btn.click(fn=manager.calculate_perplexity, outputs=perplexity_output)
+
+    return demo
+
+
+if __name__ == "__main__":
+    demo = create_interface()
+    demo.launch()

utils.py

@@ -0,0 +1,94 @@
+from typing import Final
+
+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"))
+
+checkpoints = ["HuggingFaceTB/SmolLM2-135M"]  # Inputs
+
+
+def load_model(checkpoints: list[str]) -> dict:
+    tokenizers = [
+        AutoTokenizer.from_pretrained(checkpoint) for checkpoint in checkpoints
+    ]
+
+    models = [
+        AutoModelForCausalLM.from_pretrained(
+            checkpoint,
+            device_map="auto",
+            torch_dtype=torch.bfloat16,
+        )
+        .to(dev)
+        .eval()
+        for checkpoint 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)
+    }
+
+
+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] = 1500  # Sample size for perplexity calculation
+sample_length: Final[int] = 128  # Maximum length of text to consider for perplexity
+
+loaded = load_model(checkpoints)
+
+
+def log_perplexity() -> dict:
+    # Initialize a dictionary to store perplexity
+    ppls = {checkpoint: [] for checkpoint in loaded.keys()}
+    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()}
+
+    # Calculate the mean perplexity for each model
+    mean_log_ppls: dict = {
+        checkpoint: np.mean(ppl) for checkpoint, ppl in log_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()
+    }
+
+    return {
+        "ppls": ppls,
+        "mean_ppls": mean_log_ppls,
+        "std_ppls": std_log_ppls,
+    }