diff --git a/src/gabriel/api.py b/src/gabriel/api.py index d1077df..51bf718 100644 --- a/src/gabriel/api.py +++ b/src/gabriel/api.py @@ -1133,6 +1133,8 @@ async def paraphrase( reset_files: bool = False, reasoning_effort: Optional[str] = None, reasoning_summary: Optional[str] = None, + modality: str = "text", + search_context_size: str = "medium", n_rounds: int = 1, recursive_validation: Optional[bool] = None, n_initial_candidates: int = 1, @@ -1171,6 +1173,10 @@ async def paraphrase( Whether to request JSON responses. web_search: Enable web search augmentation when supported by the model. + modality: + Modality of inputs (text, image, audio, pdf, entity, web). + search_context_size: + Web search context size when ``modality="web"``. n_parallels: Maximum concurrent paraphrase calls. use_dummy: @@ -1220,6 +1226,8 @@ async def paraphrase( use_dummy=use_dummy, reasoning_effort=reasoning_effort, reasoning_summary=reasoning_summary, + modality=modality, + search_context_size=search_context_size, n_rounds=n_rounds, recursive_validation=recursive_validation, n_initial_candidates=n_initial_candidates, diff --git a/src/gabriel/prompts/paraphrase_prompt.jinja2 b/src/gabriel/prompts/paraphrase_prompt.jinja2 index 820a9a0..36408e2 100644 --- a/src/gabriel/prompts/paraphrase_prompt.jinja2 +++ b/src/gabriel/prompts/paraphrase_prompt.jinja2 @@ -1,17 +1,24 @@ -Consider the following text: -BEGIN TEXT -{{ text }} -END TEXT -Your task is to rewrite the above text in accordance with the following instructions: -BEGIN INSTRUCTIONS -{{ instructions }} -END INSTRUCTIONS -The text should be rewritten in a way that is consistent with the instructions. -Be very diligent about following the instructions. -Output only your revised text, nothing else, no preamble, no postamble, no explanation, no nothing other than the revised text. -Consider the whole text, and follow the instructions carefully and completely. -Be meticulous and exacting; your output of rewritten text must be as faithful as possible in following the instructions. -Once again, the instructions are: +{% import "snippets.jinja2" as snip %} +{{ snip.single_entry(modality, text) }} + +Task: rewrite the provided content according to the instructions below. Be diligent, thorough, and faithful to the instructions in all cases. + BEGIN INSTRUCTIONS {{ instructions }} END INSTRUCTIONS + +Modality guidance: +- text: paraphrase the full passage, preserving meaning and key details while applying the instructions. +- entity: rewrite the entity name in a consistent way per the instructions. +- web: paraphrase/consolidate the web content found for the search term per the instructions. +- image/audio/pdf: paraphrase the observed content into text in accordance with the instructions. + +Output format (template): +{% if json_mode %} +Always output JSON only in the indicated format. +{"revised_text": ""} +{% else %} + +{% endif %} + +No preamble, no explanation, only the desired revised text alone, faithful to the provided content and in accordance with the instructions. diff --git a/src/gabriel/tasks/paraphrase.py b/src/gabriel/tasks/paraphrase.py index ed74eef..1dcce6d 100644 --- a/src/gabriel/tasks/paraphrase.py +++ b/src/gabriel/tasks/paraphrase.py @@ -12,6 +12,12 @@ from ..core.prompt_template import PromptTemplate, resolve_template from ..utils.openai_utils import get_all_responses +from ..utils import ( + load_audio_inputs, + load_image_inputs, + load_pdf_inputs, + warn_if_modality_mismatch, +) from ..utils.logging import announce_prompt_rendering # Import classifier utilities for recursive validation. Importing from @@ -50,6 +56,10 @@ class ParaphraseConfig: # augmentation. ``None`` defers to downstream defaults while ``True`` and # ``False`` explicitly enable or disable the feature. web_search: Optional[bool] = None + # What kind of input is being paraphrased (text, image, audio, pdf, entity, web). + modality: str = "text" + # Web search context size when modality is ``web``. + search_context_size: str = "medium" # Maximum number of parallel requests that will be sent to the # underlying API. Note that classification and paraphrasing share # this value for simplicity. @@ -155,7 +165,9 @@ async def run( # Convert the target column into a list of strings. We coerce # values to strings so that non-string columns (e.g. numbers) are # handled gracefully. - texts: List[str] = df_proc[column_name].astype(str).tolist() + values = df_proc[column_name].tolist() + texts: List[str] = [str(v) for v in values] + warn_if_modality_mismatch(values, self.cfg.modality, column_name=column_name) # Determine the base name for the revised column(s). base_col = self.cfg.revised_column_name or f"{column_name}_revised" # Determine how many paraphrases to produce per passage. A value @@ -190,19 +202,67 @@ async def run( prompts: List[str] = [] identifiers: List[str] = [] for idx, text in enumerate(texts): + prompt_text = ( + text + if self.cfg.modality in {"text", "entity", "web"} + else "" + ) for j in range(1, n + 1): prompts.append( - self.template.render(text=text, instructions=self.cfg.instructions) + self.template.render( + text=prompt_text, + instructions=self.cfg.instructions, + modality=self.cfg.modality, + json_mode=self.cfg.json_mode, + ) ) identifiers.append(f"row_{idx}_rev{j}") save_path = os.path.join(self.cfg.save_dir, self.cfg.file_name) + prompt_images: Optional[Dict[str, List[str]]] = None + prompt_audio: Optional[Dict[str, List[Dict[str, str]]]] = None + prompt_pdfs: Optional[Dict[str, List[Dict[str, str]]]] = None + + if self.cfg.modality == "image": + tmp: Dict[str, List[str]] = {} + for idx, val in enumerate(values): + imgs = load_image_inputs(val) + if imgs: + for j in range(1, n + 1): + tmp[f"row_{idx}_rev{j}"] = imgs + prompt_images = tmp or None + elif self.cfg.modality == "audio": + tmp_a: Dict[str, List[Dict[str, str]]] = {} + for idx, val in enumerate(values): + auds = load_audio_inputs(val) + if auds: + for j in range(1, n + 1): + tmp_a[f"row_{idx}_rev{j}"] = auds + prompt_audio = tmp_a or None + elif self.cfg.modality == "pdf": + tmp_p: Dict[str, List[Dict[str, str]]] = {} + for idx, val in enumerate(values): + pdfs = load_pdf_inputs(val) + if pdfs: + for j in range(1, n + 1): + tmp_p[f"row_{idx}_rev{j}"] = pdfs + prompt_pdfs = tmp_p or None + + web_search = ( + self.cfg.web_search + if self.cfg.web_search is not None + else self.cfg.modality == "web" + ) + kwargs.setdefault("web_search", web_search) + kwargs.setdefault("search_context_size", self.cfg.search_context_size) resp_df = await get_all_responses( prompts=prompts, identifiers=identifiers, + prompt_images=prompt_images, + prompt_audio=prompt_audio, + prompt_pdfs=prompt_pdfs, save_path=save_path, model=self.cfg.model, json_mode=self.cfg.json_mode, - web_search=self.cfg.web_search, n_parallels=self.cfg.n_parallels, use_dummy=self.cfg.use_dummy, reset_files=reset_files, @@ -227,6 +287,7 @@ async def run( resp_map: Dict[Tuple[int, int], str] = {} await self._recursive_validate( texts, + values, resp_map, approval_map, reset_files=reset_files, @@ -261,6 +322,7 @@ async def run( async def _recursive_validate( self, original_texts: List[str], + original_values: List[Any], resp_map: Dict[Tuple[int, int], str], approval_map: Dict[Tuple[int, int], bool], *, @@ -358,13 +420,27 @@ async def _recursive_validate( # exists for this key, use that paraphrase as the base # for regeneration. Otherwise, continue to use the # original. - if round_number > 0 and self.cfg.use_modified_source and key in resp_map: + if ( + round_number > 0 + and self.cfg.use_modified_source + and self.cfg.modality in {"text", "entity", "web"} + and key in resp_map + ): base_text = resp_map[key] else: - base_text = original_texts[row_idx] + base_text = ( + original_texts[row_idx] + if self.cfg.modality in {"text", "entity", "web"} + else "" + ) for cand_idx in range(candidates_per_key): prompts.append( - self.template.render(text=base_text, instructions=self.cfg.instructions) + self.template.render( + text=base_text, + instructions=self.cfg.instructions, + modality=self.cfg.modality, + json_mode=self.cfg.json_mode, + ) ) # Encode row, revision, round and candidate index in # the identifier. Revision numbers are stored one- @@ -389,13 +465,60 @@ async def _recursive_validate( self.cfg.save_dir, f"{os.path.splitext(self.cfg.file_name)[0]}_round{round_number}.csv", ) + prompt_images: Optional[Dict[str, List[str]]] = None + prompt_audio: Optional[Dict[str, List[Dict[str, str]]]] = None + prompt_pdfs: Optional[Dict[str, List[Dict[str, str]]]] = None + + if self.cfg.modality == "image": + tmp: Dict[str, List[str]] = {} + for key in to_check: + row_idx, rev_idx = key + imgs = load_image_inputs(original_values[row_idx]) + if imgs: + for cand_idx in range(candidates_per_key): + tmp[ + f"row_{row_idx}_rev{rev_idx + 1}_round{round_number}_cand{cand_idx}" + ] = imgs + prompt_images = tmp or None + elif self.cfg.modality == "audio": + tmp_a: Dict[str, List[Dict[str, str]]] = {} + for key in to_check: + row_idx, rev_idx = key + auds = load_audio_inputs(original_values[row_idx]) + if auds: + for cand_idx in range(candidates_per_key): + tmp_a[ + f"row_{row_idx}_rev{rev_idx + 1}_round{round_number}_cand{cand_idx}" + ] = auds + prompt_audio = tmp_a or None + elif self.cfg.modality == "pdf": + tmp_p: Dict[str, List[Dict[str, str]]] = {} + for key in to_check: + row_idx, rev_idx = key + pdfs = load_pdf_inputs(original_values[row_idx]) + if pdfs: + for cand_idx in range(candidates_per_key): + tmp_p[ + f"row_{row_idx}_rev{rev_idx + 1}_round{round_number}_cand{cand_idx}" + ] = pdfs + prompt_pdfs = tmp_p or None + + web_search = ( + self.cfg.web_search + if self.cfg.web_search is not None + else self.cfg.modality == "web" + ) new_resp_df = await get_all_responses( prompts=prompts, identifiers=identifiers, + prompt_images=prompt_images, + prompt_audio=prompt_audio, + prompt_pdfs=prompt_pdfs, save_path=tmp_save_path, model=self.cfg.model, json_mode=self.cfg.json_mode, - web_search=self.cfg.web_search, + web_search=web_search, + search_context_size=self.cfg.search_context_size, n_parallels=self.cfg.n_parallels, use_dummy=self.cfg.use_dummy, reset_files=reset_files, diff --git a/tests/test_basic.py b/tests/test_basic.py index e7c748a..711cf77 100644 --- a/tests/test_basic.py +++ b/tests/test_basic.py @@ -1582,6 +1582,73 @@ def test_paraphrase_api(tmp_path): assert "txt_revised_1" in df_multi.columns and "txt_revised_2" in df_multi.columns +def test_paraphrase_modalities_forward_media(monkeypatch, tmp_path): + captured = {} + + async def fake_get_all_responses(*, prompts, identifiers, **kwargs): + captured["prompts"] = list(prompts) + captured["identifiers"] = list(identifiers) + captured["prompt_images"] = kwargs.get("prompt_images") + captured["prompt_audio"] = kwargs.get("prompt_audio") + captured["prompt_pdfs"] = kwargs.get("prompt_pdfs") + captured["web_search"] = kwargs.get("web_search") + captured["search_context_size"] = kwargs.get("search_context_size") + return pd.DataFrame({"Identifier": identifiers, "Response": prompts}) + + monkeypatch.setattr("gabriel.tasks.paraphrase.get_all_responses", fake_get_all_responses) + + df_image = pd.DataFrame({"media": ["data:image/png;base64,xyz"]}) + asyncio.run( + gabriel.paraphrase( + df_image, + "media", + instructions="caption it", + save_dir=str(tmp_path / "para_image"), + modality="image", + ) + ) + assert captured["prompt_images"]["row_0_rev1"] == ["data:image/png;base64,xyz"] + + df_audio = pd.DataFrame({"media": [{"format": "mp3", "data": "abcd"}]}) + asyncio.run( + gabriel.paraphrase( + df_audio, + "media", + instructions="summarize it", + save_dir=str(tmp_path / "para_audio"), + modality="audio", + ) + ) + assert captured["prompt_audio"]["row_0_rev1"][0]["format"] == "mp3" + + df_pdf = pd.DataFrame({"media": ["data:application/pdf;base64,abcd"]}) + asyncio.run( + gabriel.paraphrase( + df_pdf, + "media", + instructions="rewrite", + save_dir=str(tmp_path / "para_pdf"), + modality="pdf", + ) + ) + assert captured["prompt_pdfs"]["row_0_rev1"][0]["file_data"].startswith( + "data:application/pdf" + ) + + df_web = pd.DataFrame({"query": ["Mount Fuji"]}) + asyncio.run( + gabriel.paraphrase( + df_web, + "query", + instructions="summarize", + save_dir=str(tmp_path / "para_web"), + modality="web", + ) + ) + assert captured["web_search"] is True + assert captured["search_context_size"] == "medium" + + def test_paraphrase_n_rounds_not_forwarded(monkeypatch, tmp_path): captured = {} @@ -1607,7 +1674,16 @@ async def fake_get_all_responses(*, prompts, identifiers, **kwargs): def test_paraphrase_approval_column_recursive(monkeypatch, tmp_path): - async def fake_recursive_validate(self, original_texts, resp_map, approval_map, *, reset_files, max_rounds): + async def fake_recursive_validate( + self, + original_texts, + original_values, + resp_map, + approval_map, + *, + reset_files, + max_rounds, + ): for idx in range(len(original_texts)): resp_map[(idx, 0)] = f"approved {idx}" approval_map[(idx, 0)] = (idx % 2 == 0)