Create fitness.py

backend/fitness.py

@@ -0,0 +1,74 @@
+import subprocess
+import tempfile
+import os
+
+def evaluate_model(model_path: str, dataset_path: str, custom_script: str = None) -> float:
+    """
+    Default fitness: perplexity for causal LM, or CLIPScore for diffusion.
+    Custom script can override by running a Python file that defines fitness(model_path, dataset_path) -> float.
+    """
+    if custom_script:
+        # Write script to temp file and execute
+        with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
+            f.write(custom_script)
+            script_path = f.name
+        try:
+            result = subprocess.run(
+                ["python", script_path, model_path, dataset_path],
+                capture_output=True, text=True, timeout=300
+            )
+            if result.returncode != 0:
+                raise RuntimeError(f"Fitness script error: {result.stderr}")
+            return float(result.stdout.strip())
+        finally:
+            os.unlink(script_path)
+    else:
+        # Auto-detect model type
+        from transformers import AutoConfig
+        config = AutoConfig.from_pretrained(model_path)
+        if hasattr(config, "architectures") and any("LM" in a for a in config.architectures):
+            return compute_perplexity(model_path, dataset_path)
+        else:
+            return compute_clip_score(model_path, dataset_path)
+
+def compute_perplexity(model_path, dataset_path):
+    import torch
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float32, low_cpu_mem_usage=True)
+    model.eval()
+    total_loss = 0
+    total_tokens = 0
+    with open(dataset_path, "r") as f:
+        text = f.read()
+    encodings = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
+    with torch.no_grad():
+        outputs = model(**encodings, labels=encodings["input_ids"])
+        loss = outputs.loss
+        return loss.item()
+    return float("inf")
+
+def compute_clip_score(model_path, dataset_path):
+    # For diffusion, generate images from prompts in dataset file and compute CLIP score
+    import torch
+    from transformers import CLIPProcessor, CLIPModel
+    from diffusers import StableDiffusionPipeline
+    import numpy as np
+    from PIL import Image
+    pipe = StableDiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float32, safety_checker=None)
+    clip = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
+    processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+    prompts = []
+    with open(dataset_path, "r") as f:
+        for line in f:
+            if line.strip():
+                prompts.append(line.strip())
+    images = []
+    for prompt in prompts[:5]:  # limit for speed
+        img = pipe(prompt, num_inference_steps=10).images[0]
+        images.append(img)
+    inputs = processor(text=prompts[:5], images=images, return_tensors="pt", padding=True)
+    outputs = clip(**inputs)
+    logits_per_image = outputs.logits_per_image
+    score = logits_per_image.diag().mean().item()
+    return -score  # lower is better (closer to 1)