JairoDanielMT's picture
Upload folder using huggingface_hub
4ef6c2b verified
Raw
History Blame Contribute Delete
6.87 kB
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/")