| import gradio as gr |
|
|
| from utils import load_model, log_perplexity |
|
|
|
|
| class Manager: |
| """Class to manage model loading and perplexity calculation state.""" |
|
|
| def __init__(self): |
| self.loaded = None |
|
|
| def load_models(self, checkpoint_input_str: str) -> str: |
| """Load models from a comma-separated string of checkpoint names.""" |
| checkpoints = [ |
| ckpt.strip() for ckpt in checkpoint_input_str.split(",") if ckpt.strip() |
| ] |
|
|
| if not checkpoints: |
| return "Please enter at least one model checkpoint name." |
|
|
| try: |
| self.loaded = load_model(checkpoints) |
| return "Models loaded successfully!" |
| except Exception as e: |
| return f"Model loading failed: {e}" |
|
|
| def perplexity( |
| self, |
| num_samples: int | None = None, |
| sample_length: int | None = None, |
| ) -> dict | str: |
| """Calculate perplexity using the loaded models.""" |
| 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: |
| return log_perplexity(self.loaded, num_samples, sample_length) |
| except Exception as e: |
| return f"Perplexity calculation failed: {e}" |
|
|
|
|
| def make_interface() -> gr.Blocks: |
| """Create and return the Gradio interface.""" |
| manager = Manager() |
|
|
| with gr.Blocks() as demo: |
| gr.Markdown("# LLM PPLs") |
|
|
| checkpoints = gr.Textbox( |
| label="Checkpoints", value="HuggingFaceTB/SmolLM2-135M" |
| ) |
|
|
| load_btn = gr.Button("Load Models", variant="primary") |
|
|
| 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") |
|
|
| |
| load_btn.click( |
| fn=manager.load_models, |
| inputs=checkpoints, |
| outputs=load_output, |
| ) |
|
|
| perplexity_btn.click( |
| fn=manager.perplexity, |
| inputs=[num_samples, sample_length], |
| outputs=perplexity_output, |
| ) |
|
|
| return demo |
|
|
|
|
| if __name__ == "__main__": |
| demo = make_interface() |
| demo.launch() |
|
|
