"""Data preprocessing pipeline for AskBeforeAnswer. This module is responsible for loading the raw datasets and processing them into the structured JSONL formats required for Supervised Fine-Tuning (SFT) and Direct Preference Optimization (DPO). """ import ast import json import logging import re from typing import Any, Dict, List, Optional import pandas as pd from datasets import load_dataset from omegaconf import DictConfig from tqdm import tqdm logger = logging.getLogger(__name__) def clean_facets(facets: Any) -> List[str]: """Clean and parse facet lists.""" if isinstance(facets, list): return facets if isinstance(facets, str): try: val = ast.literal_eval(facets) return val if isinstance(val, list) else [] except Exception: return [] return [] def clean_response(resp: Any) -> str: """Clean model responses.""" if resp is None: return "" if isinstance(resp, list): return str(resp[0]) if resp else "" if isinstance(resp, str) and resp.startswith("[") and resp.endswith("]"): try: lst = ast.literal_eval(resp) return str(lst[0]) if lst else "" except Exception: pass return str(resp).strip() class SyntheticGenerator: """Generates synthetic reasoning, facets, and responses using a lightweight LLM.""" SYSTEM_PROMPT = """You are a clarification-seeking question-understanding agent. Your job: 1. Determine whether the user question is AMBIGUOUS. 2. If ambiguous -> identify the facets of ambiguity and ask a clarifying question. 3. If unambiguous -> answer directly. 4. ALWAYS output in the EXACT format below (no deviations): Action: Clarify | Answer Reasoning: Facets: [list, of, facets] # empty list [] if the question is unambiguous Response: -------------------- ### FEW-SHOT EXAMPLES -------------------- Example 1: User Question: "Who founded Apple?" Action: Clarify Reasoning: The question refers to "Apple" but multiple founders exist (Steve Jobs, Steve Wozniak, Ronald Wayne). Clarification is needed to know which person the user is asking about. Facets: ["Which founder?", "Steve Jobs vs Steve Wozniak vs Ronald Wayne"] Response: Do you mean Steve Jobs, Steve Wozniak, or Ronald Wayne? Example 2: User Question: "What is 2+2?" Action: Answer Reasoning: The question is clear, numeric, and has only one interpretation. No facets of ambiguity exist. Facets: [] Response: 4 -------------------- Follow the format STRICTLY for every response. """ def __init__(self, model_id: str, batch_size: int = 8, max_new_tokens: int = 256): import torch from transformers import AutoModelForCausalLM, AutoTokenizer from transformers import logging as tf_logging tf_logging.set_verbosity_error() logger.info(f"Loading synthetic generator model: {model_id}") self.tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) # Handle pad token and left-padding for batched generation if self.tokenizer.pad_token is None: self.tokenizer.pad_token = self.tokenizer.eos_token self.tokenizer.padding_side = "left" kwargs = {"device_map": "auto"} if torch.cuda.is_available(): kwargs["torch_dtype"] = torch.bfloat16 elif torch.backends.mps.is_available(): kwargs["torch_dtype"] = torch.float16 self.model = AutoModelForCausalLM.from_pretrained( model_id, trust_remote_code=True, **kwargs ) self.model.eval() self.batch_size = batch_size self.max_new_tokens = max_new_tokens def _parse_output(self, text: str) -> Dict[str, Any]: """Parses the generated text into structured fields.""" action_match = re.search(r"Action:\s*(Clarify|Answer)", text, re.IGNORECASE) action = action_match.group(1).title() if action_match else "Answer" reasoning_match = re.search( r"Reasoning:\s*(.*?)(?=\nFacets:|\Z)", text, re.DOTALL | re.IGNORECASE ) reasoning = ( reasoning_match.group(1).strip() if reasoning_match else "The question is missing specific details." ) facets_match = re.search( r"Facets:\s*(\[.*?\])", text, re.DOTALL | re.IGNORECASE ) facets = [] if facets_match: try: facets = ast.literal_eval(facets_match.group(1).strip()) except (SyntaxError, ValueError): pass response_match = re.search(r"Response:\s*(.*)", text, re.DOTALL | re.IGNORECASE) response = response_match.group(1).strip() if response_match else "" return { "action": action, "reasoning": reasoning, "facets": facets if isinstance(facets, list) else [], "response": response, } def generate_batch(self, questions: List[str]) -> List[Dict[str, Any]]: """Generates structured outputs for a batch of questions.""" import torch prompts = [ self.tokenizer.apply_chat_template( [ {"role": "system", "content": self.SYSTEM_PROMPT}, {"role": "user", "content": q}, ], tokenize=False, add_generation_prompt=True, ) for q in questions ] inputs = self.tokenizer(prompts, return_tensors="pt", padding=True).to( self.model.device ) with torch.no_grad(): outputs = self.model.generate( **inputs, max_new_tokens=self.max_new_tokens, do_sample=False, temperature=None, top_p=None, top_k=None, pad_token_id=self.tokenizer.pad_token_id, eos_token_id=self.tokenizer.eos_token_id, ) results = [] for i, output in enumerate(outputs): input_len = inputs.input_ids[i].shape[0] generated_text = self.tokenizer.decode( output[input_len:], skip_special_tokens=True ) results.append(self._parse_output(generated_text)) return results def extract_qa_data( dataset_name: str, split: str = "train", max_samples: Optional[int] = None, synthetic_cfg: Optional[DictConfig] = None, ) -> pd.DataFrame: """Load AmbigQA dataset and extract question and annotation details. Args: dataset_name (str): The name or path of the dataset to load from HuggingFace. split (str): The dataset split to process (e.g., 'train', 'validation'). max_samples (Optional[int]): Maximum number of rows to extract. synthetic_cfg (Optional[DictConfig]): Configuration for LLM synthetic generation. Returns: pd.DataFrame: A DataFrame containing the extracted questions and facets. """ logger.info(f"Loading dataset {dataset_name} ({split})...") ds = load_dataset(dataset_name, split=split) if max_samples: ds = ds.select(range(min(len(ds), max_samples))) rows = [] questions = [] # Base extraction for entry in ds: ann = entry["annotations"] ann_type = ann["type"][0] if isinstance(ann["type"], list) else ann["type"] qa_pairs = ann.get("qaPairs", []) if len(qa_pairs) == 1 and isinstance(qa_pairs[0], list): qa_pairs = qa_pairs[0] single_ans_list = ann.get("answer", []) single_ans_flat = [] if isinstance(single_ans_list, list): for ans in single_ans_list: if isinstance(ans, list): single_ans_flat.extend(ans) elif isinstance(ans, str): single_ans_flat.append(ans) is_ambiguous = ann_type == "multipleQAs" questions.append(entry["question"]) # Determine base truth for fallback and correctness checks base_action = "Clarify" if is_ambiguous else "Answer" base_pos_resp = ( qa_pairs[0].get("question", "Could you clarify?") if is_ambiguous and qa_pairs else (single_ans_flat[0] if single_ans_flat else "Direct answer.") ) rows.append( { "question": entry["question"], "is_ambiguous": is_ambiguous, "action": base_action, "facets": ["Entity Reference"] if is_ambiguous else [], "reasoning": ( "The question is missing specific details." if is_ambiguous else "The question is clear." ), "positive_response": base_pos_resp, } ) df = pd.DataFrame(rows) # Apply synthetic generation if enabled if synthetic_cfg and synthetic_cfg.get("enabled", False): generator = SyntheticGenerator( model_id=synthetic_cfg.get("model_id", "Qwen/Qwen2.5-3B-Instruct"), batch_size=synthetic_cfg.get("batch_size", 8), max_new_tokens=synthetic_cfg.get("max_new_tokens", 256), ) logger.info( f"Generating synthetic annotations for {len(questions)} examples..." ) batch_size = generator.batch_size for i in tqdm(range(0, len(questions), batch_size), desc="Synthetic Gen"): batch_q = questions[i : i + batch_size] batch_results = generator.generate_batch(batch_q) for j, res in enumerate(batch_results): idx = i + j # Overwrite placeholder fields with LLM-generated high-quality data # We enforce the ground-truth action from AmbigNQ # to ensure correct labels is_amb = df.at[idx, "is_ambiguous"] df.at[idx, "reasoning"] = res["reasoning"] df.at[idx, "facets"] = res["facets"] if is_amb else [] # If the LLM successfully generated a response that matches # the target action, use it # Otherwise, fallback to the dataset's ground truth # to avoid noisy positives if res["action"] == df.at[idx, "action"] and res["response"]: df.at[idx, "positive_response"] = res["response"] return df def prepare_sft_dataset(df: pd.DataFrame, output_path: str) -> None: """Format DataFrame into SFT JSONL format and save to disk.""" logger.info(f"Preparing SFT dataset to {output_path}...") records = [] system_instruction = ( "You are a helpful assistant. " "Given a question, you must decide whether it is ambiguous or not. " "Output MUST follow this format:\n" "Action: Clarify|Answer\n" "Reasoning: \n" "Facets: \n" "Response: " ) for _, row in df.iterrows(): is_amb = row["is_ambiguous"] facets = clean_facets(row["facets"]) records.append( { "instruction": system_instruction, "input": row["question"], "output": { "action": row["action"], "reasoning": row["reasoning"], "facets": facets if is_amb else [], "response": clean_response(row["positive_response"]), }, } ) with open(output_path, "w") as f: for r in records: f.write(json.dumps(r) + "\n") logger.info(f"Saved {len(records)} records for SFT.") def prepare_dpo_dataset(df: pd.DataFrame, output_path: str) -> None: """Format DataFrame into DPO JSONL format and save to disk.""" logger.info(f"Preparing DPO dataset to {output_path}...") records = [] for _, row in df.iterrows(): # DPO requires chosen vs rejected action = row["action"] reasoning = row["reasoning"] is_amb = row["is_ambiguous"] facets_list = clean_facets(row["facets"]) if is_amb else [] facets_str = str(facets_list) chosen_resp = clean_response(row["positive_response"]) chosen = ( f"Action: {action}\n" f"Reasoning: {reasoning}\n" f"Facets: {facets_str}\n" f"Response: {chosen_resp}" ) rejected_action = "Answer" if action == "Clarify" else "Clarify" rejected = ( f"Action: {rejected_action}\n" "Reasoning: Incorrect reasoning: the model misunderstood the question.\n" 'Facets: ["Incorrect Interpretation"]\n' "Response: " + ("I don't know." if action == "Clarify" else "Could you clarify?") ) records.append( { "prompt": row["question"], "chosen": chosen, "rejected": rejected, } ) with open(output_path, "w") as f: for r in records: f.write(json.dumps(r) + "\n") logger.info(f"Saved {len(records)} records for DPO.")