Spaces:
Sleeping
Sleeping
| import json | |
| import random | |
| import os | |
| from typing import List, Dict, Any | |
| from src.parser.parser import Parser | |
| from src.ontology.matcher import ConceptMatcher | |
| from src.embeddings.engine import EmbeddingEngine | |
| from src.enrichment.enricher import Enricher | |
| class BenchmarkSuite: | |
| def __init__(self, ontology_dir: str): | |
| self.ontology_dir = ontology_dir | |
| self.concepts = self._load_all_concepts() | |
| def _load_all_concepts(self) -> Dict[str, List[str]]: | |
| concepts = {} | |
| for filename in os.listdir(self.ontology_dir): | |
| if filename.endswith(".json"): | |
| path = os.path.join(self.ontology_dir, filename) | |
| category = filename.replace(".json", "") | |
| with open(path, "r", encoding="utf-8") as f: | |
| data = json.load(f) | |
| concepts[category] = [item["canonical"] for item in data] | |
| return concepts | |
| def generate_benchmarks(self, count: int = 1000) -> List[Dict[str, Any]]: | |
| benchmarks = [] | |
| # Core characters | |
| sonic_chars = ["Sonic", "Amy Rose", "Shadow", "Rouge", "Blaze", "Silver", "Knuckles", "Tails", "Cream"] | |
| for _ in range(count): | |
| # Mix strategies | |
| strategy = random.random() | |
| if strategy < 0.3: # Character focus | |
| char = random.choice(sonic_chars) | |
| clothing = random.choice(self.concepts.get("clothing", ["dress"])) | |
| scene = random.choice(self.concepts.get("scenes", ["beach"])) | |
| input_text = f"{char} wearing a {clothing} at the {scene}" | |
| expected = [char, clothing, scene] | |
| elif strategy < 0.6: # Attribute focus | |
| style = random.choice(self.concepts.get("styles", ["anime"])) | |
| hair = random.choice(self.concepts.get("hair_colors", ["blue hair"])) | |
| eyes = random.choice(self.concepts.get("eye_colors", ["green eyes"])) | |
| input_text = f"{style} style girl with {hair} and {eyes}" | |
| expected = [style, hair, eyes] | |
| else: # Random mix | |
| parts = [] | |
| expected = [] | |
| for cat in ["clothing", "accessories", "scenes", "emotions", "lighting"]: | |
| if random.random() > 0.5 and self.concepts.get(cat): | |
| val = random.choice(self.concepts[cat]) | |
| parts.append(val) | |
| expected.append(val) | |
| input_text = " ".join(parts) | |
| if input_text: | |
| benchmarks.append({ | |
| "input": input_text, | |
| "expected_entities": expected | |
| }) | |
| return benchmarks | |
| def evaluate(self, benchmarks: List[Dict[str, Any]], parser: Parser, enricher: Enricher) -> Dict[str, Any]: | |
| results = [] | |
| total_precision = 0 | |
| total_recall = 0 | |
| exact_match_count = 0 | |
| entity_resolved_count = 0 | |
| total_expected_entities = 0 | |
| for case in benchmarks: | |
| ir = parser.parse(case["input"]) | |
| ir = enricher.enrich(ir) | |
| # Extract all found canonicals | |
| found = set() | |
| for char in ir.characters: | |
| found.add(char.name) | |
| found.update(char.appearance) | |
| found.update(char.clothing) | |
| found.update(char.accessories) | |
| found.update(char.pose) | |
| found.update(char.expression) | |
| found.update(ir.scene.locations) | |
| found.update(ir.scene.lighting) | |
| found.update(ir.scene.atmosphere) | |
| found.update(ir.style) | |
| found.update(ir.effects) | |
| expected = set(case["expected_entities"]) | |
| intersection = found.intersection(expected) | |
| precision = len(intersection) / len(found) if found else 0 | |
| recall = len(intersection) / len(expected) if expected else 1.0 | |
| total_precision += precision | |
| total_recall += recall | |
| total_expected_entities += len(expected) | |
| entity_resolved_count += len(intersection) | |
| if recall == 1.0: | |
| exact_match_count += 1 | |
| results.append({ | |
| "input": case["input"], | |
| "expected": list(expected), | |
| "found": list(found), | |
| "precision": precision, | |
| "recall": recall | |
| }) | |
| avg_precision = total_precision / len(benchmarks) | |
| avg_recall = total_recall / len(benchmarks) | |
| f1 = 2 * (avg_precision * avg_recall) / (avg_precision + avg_recall) if (avg_precision + avg_recall) else 0 | |
| return { | |
| "metrics": { | |
| "precision": round(avg_precision, 4), | |
| "recall": round(avg_recall, 4), | |
| "f1": round(f1, 4), | |
| "exact_match_rate": round(exact_match_count / len(benchmarks), 4), | |
| "entity_resolution_accuracy": round(entity_resolved_count / total_expected_entities if total_expected_entities else 0, 4) | |
| }, | |
| "total_cases": len(benchmarks), | |
| "results": results | |
| } | |
| if __name__ == "__main__": | |
| suite = BenchmarkSuite("data/ontology") | |
| print("Generating 1000 benchmarks...") | |
| benchmarks = suite.generate_benchmarks(1000) | |
| with open("benchmarks/prompts.json", "w", encoding="utf-8") as f: | |
| json.dump(benchmarks, f, indent=2) | |
| print("Running evaluation...") | |
| matcher = ConceptMatcher("data/ontology") | |
| engine = EmbeddingEngine(index_dir="data/faiss_indices") | |
| engine.load_index() | |
| parser = Parser(matcher, engine) | |
| enricher = Enricher(matcher) | |
| report = suite.evaluate(benchmarks, parser, enricher) | |
| os.makedirs("reports", exist_ok=True) | |
| with open("reports/benchmark_results.json", "w", encoding="utf-8") as f: | |
| json.dump(report, f, indent=2) | |
| with open("reports/benchmark_results.md", "w", encoding="utf-8") as f: | |
| f.write("# Benchmark Results\n\n") | |
| m = report["metrics"] | |
| f.write(f"- **Total Cases**: {report['total_cases']}\n") | |
| f.write(f"- **Precision**: {m['precision']}\n") | |
| f.write(f"- **Recall**: {m['recall']}\n") | |
| f.write(f"- **F1 Score**: {m['f1']}\n") | |
| f.write(f"- **Exact Match Rate**: {m['exact_match_rate']}\n") | |
| f.write(f"- **Entity Resolution Accuracy**: {m['entity_resolution_accuracy']}\n") | |
| print("Benchmark evaluation complete. Reports generated in reports/") | |