| 30 | |
| 31 | |
| 32 | class JsonModeEval(Evaluator): |
| 33 | |
| 34 | def __init__(self, |
| 35 | dataset_path: Optional[str] = None, |
| 36 | num_samples: Optional[int] = None, |
| 37 | random_seed: int = 0, |
| 38 | apply_chat_template: bool = True, |
| 39 | system_prompt: Optional[str] = None): |
| 40 | if not apply_chat_template: |
| 41 | raise ValueError( |
| 42 | f"{self.__class__.__name__} requires apply_chat_template=True.") |
| 43 | super().__init__(random_seed=random_seed, |
| 44 | apply_chat_template=apply_chat_template, |
| 45 | system_prompt=system_prompt) |
| 46 | if dataset_path is None: |
| 47 | dataset_path = "NousResearch/json-mode-eval" |
| 48 | self.data = datasets.load_dataset(dataset_path, |
| 49 | split="train", |
| 50 | trust_remote_code=True) |
| 51 | self.data = self.data.shuffle(random_seed) |
| 52 | if num_samples is None: |
| 53 | self.num_samples = self.data.num_rows |
| 54 | else: |
| 55 | self.num_samples = min(num_samples, self.data.num_rows) |
| 56 | |
| 57 | def generate_samples(self) -> Iterable[tuple]: |
| 58 | for i, sample in enumerate(self.data): |
| 59 | if i >= self.num_samples: |
| 60 | break |
| 61 | schema = sample["schema"] |
| 62 | if os.environ.get("TRTLLM_XGUIDANCE_LENIENT") == "1": |
| 63 | schema = json.loads(schema) |
| 64 | schema["x-guidance"] = {"lenient": True} |
| 65 | schema = json.dumps(schema) |
| 66 | sampling_args = { |
| 67 | "guided_decoding": GuidedDecodingParams(json=schema) |
| 68 | } |
| 69 | yield sample["prompt"], sampling_args, sample["completion"], sample[ |
| 70 | "schema"] |
| 71 | |
| 72 | def compute_score(self, outputs: List[RequestOutput], references: List[str], |
| 73 | schemas: List[str]) -> float: |
| 74 | all_corrections, all_grammar_corrections = [], [] |
| 75 | for output, ref, schema in zip(outputs, references, schemas): |
| 76 | try: |
| 77 | output_json = json.loads(output.outputs[0].text) |
| 78 | jsonschema.validate(output_json, json.loads(schema)) |
| 79 | except (json.JSONDecodeError, jsonschema.ValidationError): |
| 80 | all_corrections.append(False) |
| 81 | all_grammar_corrections.append(False) |
| 82 | continue |
| 83 | all_corrections.append(output_json == json.loads(ref)) |
| 84 | all_grammar_corrections.append(True) |
| 85 | |
| 86 | acc = np.mean(all_corrections) * 100 |
| 87 | logger.info( |
| 88 | f"JSON Mode Eval accuracy: {acc:.2f} ({len(all_corrections)})") |
| 89 | grammar_acc = np.mean(all_grammar_corrections) * 100 |