Upload evaluate_model.py with huggingface_hub

evaluate_model.py

@@ -0,0 +1,213 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "transformers>=4.45.0",
+#     "peft>=0.7.0",
+#     "torch",
+#     "huggingface_hub",
+# ]
+# ///
+"""
+Evaluate QMD query expansion model quality.
+
+Generates expansions for test queries and outputs results for review.
+"""
+
+import json
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
+
+# Test queries covering different QMD use cases
+TEST_QUERIES = [
+    # Technical documentation
+    "how to configure authentication",
+    "typescript async await",
+    "docker compose networking",
+    "git rebase vs merge",
+    "react useEffect cleanup",
+
+    # Short/ambiguous queries
+    "auth",
+    "config",
+    "setup",
+    "api",
+
+    # Personal notes / journals style
+    "meeting notes project kickoff",
+    "ideas for new feature",
+    "todo list app architecture",
+
+    # Research / learning
+    "what is dependency injection",
+    "difference between sql and nosql",
+    "kubernetes vs docker swarm",
+
+    # Error/debugging
+    "connection timeout error",
+    "memory leak debugging",
+    "cors error fix",
+
+    # Complex queries
+    "how to implement caching with redis in nodejs",
+    "best practices for api rate limiting",
+    "setting up ci cd pipeline with github actions",
+]
+
+PROMPT_TEMPLATE = """You are a search query optimization expert. Transform the query into retrieval-optimized outputs.
+
+Query: {query}
+
+Output format:
+lex: {{keyword variation}}
+vec: {{semantic reformulation}}
+hyde: {{hypothetical document passage}}
+
+Output:"""
+
+
+def load_model(model_name: str, base_model: str = "Qwen/Qwen3-0.6B"):
+    """Load the finetuned model."""
+    print(f"Loading tokenizer from {base_model}...")
+    tokenizer = AutoTokenizer.from_pretrained(base_model)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    print(f"Loading base model...")
+    base = AutoModelForCausalLM.from_pretrained(
+        base_model,
+        torch_dtype=torch.bfloat16,
+        device_map="auto",
+    )
+
+    print(f"Loading adapter from {model_name}...")
+    model = PeftModel.from_pretrained(base, model_name)
+    model.eval()
+
+    return model, tokenizer
+
+
+def generate_expansion(model, tokenizer, query: str, max_new_tokens: int = 200) -> str:
+    """Generate query expansion."""
+    prompt = PROMPT_TEMPLATE.format(query=query)
+
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            temperature=0.7,
+            do_sample=True,
+            pad_token_id=tokenizer.pad_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+        )
+
+    # Decode and extract just the generated part
+    full_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    # Remove the prompt to get just the expansion
+    if "Output:" in full_output:
+        expansion = full_output.split("Output:")[-1].strip()
+    else:
+        expansion = full_output[len(prompt):].strip()
+
+    return expansion
+
+
+def evaluate_expansion(query: str, expansion: str) -> dict:
+    """Basic automatic evaluation metrics."""
+    lines = expansion.strip().split("\n")
+
+    has_lex = any(l.strip().startswith("lex:") for l in lines)
+    has_vec = any(l.strip().startswith("vec:") for l in lines)
+    has_hyde = any(l.strip().startswith("hyde:") for l in lines)
+
+    # Count valid lines
+    valid_lines = sum(1 for l in lines if l.strip().startswith(("lex:", "vec:", "hyde:")))
+
+    # Check for repetition
+    contents = []
+    for l in lines:
+        if ":" in l:
+            contents.append(l.split(":", 1)[1].strip().lower())
+    unique_contents = len(set(contents))
+
+    return {
+        "has_lex": has_lex,
+        "has_vec": has_vec,
+        "has_hyde": has_hyde,
+        "valid_lines": valid_lines,
+        "total_lines": len(lines),
+        "unique_contents": unique_contents,
+        "format_score": (has_lex + has_vec + has_hyde) / 3,
+    }
+
+
+def main():
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", default="tobil/qmd-query-expansion-0.6B",
+                        help="Model to evaluate")
+    parser.add_argument("--base-model", default="Qwen/Qwen3-0.6B",
+                        help="Base model")
+    parser.add_argument("--output", default="evaluation_results.json",
+                        help="Output file for results")
+    parser.add_argument("--queries", type=str, help="Custom queries file (one per line)")
+    args = parser.parse_args()
+
+    # Load custom queries if provided
+    queries = TEST_QUERIES
+    if args.queries:
+        with open(args.queries) as f:
+            queries = [l.strip() for l in f if l.strip()]
+
+    # Load model
+    model, tokenizer = load_model(args.model, args.base_model)
+
+    # Run evaluation
+    results = []
+    print(f"\n{'='*70}")
+    print("EVALUATION RESULTS")
+    print(f"{'='*70}\n")
+
+    for i, query in enumerate(queries, 1):
+        print(f"[{i}/{len(queries)}] Query: {query}")
+        print("-" * 50)
+
+        expansion = generate_expansion(model, tokenizer, query)
+        metrics = evaluate_expansion(query, expansion)
+
+        print(expansion)
+        print(f"\n  Format: {'✓' if metrics['format_score'] == 1.0 else '⚠'} "
+              f"(lex:{metrics['has_lex']}, vec:{metrics['has_vec']}, hyde:{metrics['has_hyde']})")
+        print(f"  Lines: {metrics['valid_lines']}/{metrics['total_lines']} valid, "
+              f"{metrics['unique_contents']} unique")
+        print()
+
+        results.append({
+            "query": query,
+            "expansion": expansion,
+            "metrics": metrics,
+        })
+
+    # Summary
+    print(f"\n{'='*70}")
+    print("SUMMARY")
+    print(f"{'='*70}")
+
+    avg_format = sum(r["metrics"]["format_score"] for r in results) / len(results)
+    full_format = sum(1 for r in results if r["metrics"]["format_score"] == 1.0)
+
+    print(f"  Total queries: {len(results)}")
+    print(f"  Average format score: {avg_format:.2%}")
+    print(f"  Full format compliance: {full_format}/{len(results)} ({full_format/len(results):.0%})")
+
+    # Save results
+    with open(args.output, "w") as f:
+        json.dump(results, f, indent=2)
+    print(f"\n  Results saved to: {args.output}")
+
+
+if __name__ == "__main__":
+    main()