Repository navigation
Expand file tree
/
Copy pathrefiner.py
More file actions
151 lines (138 loc) · 6.6 KB
/
Copy pathrefiner.py
File metadata and controls
151 lines (138 loc) · 6.6 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
import asyncio
import json
from typing import List, Optional
from buffdata.engine.client import GeminiClient, LLMClient, ProviderError
from buffdata.engine.limiter import AsyncRateLimiter
from buffdata.models.schemas import BatchRefinementResponse, DatasetItem, RefinementResult
REFINER_SYSTEM_PROMPT = """You are an expert AI Training Data Refiner.
Your mission is to upgrade the quality of training data for state-of-the-art LLMs.
Apply the following strict guidelines:
1. Strip AI Preamble & Boilerplate: Remove conversational filler like "Certainly! Here is...", "As an AI language model...", "I hope this helps!", etc.
2. Elevate Reasoning Depth: Where appropriate, provide clear, step-by-step chain-of-thought derivations, logical deductions, or explanations before giving final conclusions.
3. Perfect Formatting & Syntax: Fix broken markdown, ensure code blocks have appropriate language identifiers (e.g. ```python), proper indentation, and clean tables/lists.
4. Improve Prompt Precision: Disambiguate unclear instructions and eliminate meta-prompt artifacts without altering the original user intent.
5. Preserve Core Meaning: Do not change the fundamental task or introduce fabricated assertions.
"""
class DataRefiner:
"""Polishes prompts, expands reasoning steps, fixes formatting, and removes boilerplate."""
def __init__(
self,
client: Optional[LLMClient] = None,
limiter: Optional[AsyncRateLimiter] = None,
model: Optional[str] = None,
):
self.client = client or GeminiClient()
self.limiter = limiter or AsyncRateLimiter()
self.model = model or self.client.default_model
async def refine_item_async(self, item: DatasetItem, mode: str = "all") -> DatasetItem:
"""Refine prompt and/or response of a single dataset item."""
prompt, response = item.get_prompt_and_response()
refine_instruction = f"""Refine this training sample (mode: {mode}):
[ORIGINAL PROMPT]
{prompt}
[ORIGINAL RESPONSE]
{response}
"""
try:
result: RefinementResult = await self.limiter.execute_with_retry(
lambda: self.client.generate_structured_async(
prompt=refine_instruction,
response_schema=RefinementResult,
model=self.model,
system_instruction=REFINER_SYSTEM_PROMPT,
temperature=0.2,
)
)
new_prompt = result.refined_prompt if (mode in ("all", "prompt_only") and result.refined_prompt) else prompt
new_response = result.refined_response if mode in ("all", "response_only") else response
item.update_content(new_prompt=new_prompt, new_response=new_response)
item.metadata["refined"] = True
item.metadata["refinement_info"] = {
"reasoning_added": result.reasoning_added,
"formatting_fixed": result.formatting_fixed,
"artifacts_removed": result.artifacts_removed,
"explanation": result.explanation_of_changes,
}
except ProviderError:
raise
except Exception as e:
item.metadata["refinement_error"] = str(e)
return item
async def _refine_chunk(
self,
chunk: List[DatasetItem],
mode: str,
max_field_chars: int,
on_progress: Optional[callable],
) -> None:
records = []
for item in chunk:
prompt, response = item.get_prompt_and_response()
records.append({
"item_id": item.id,
"prompt": prompt[:max_field_chars],
"response": response[:max_field_chars],
})
batch_prompt = (
f"Refine every training sample below (mode: {mode}). Return exactly one entry per "
"item_id, preserve each item_id verbatim, and apply the refinement guidelines "
"independently to each record.\n\n" + json.dumps(records, ensure_ascii=False)
)
try:
response: BatchRefinementResponse = await self.limiter.execute_with_retry(
lambda: self.client.generate_structured_async(
prompt=batch_prompt,
response_schema=BatchRefinementResponse,
model=self.model,
system_instruction=REFINER_SYSTEM_PROMPT,
temperature=0.2,
)
)
by_id = {entry.item_id: entry.result for entry in response.entries}
except ProviderError:
raise
except Exception as exc:
for item in chunk:
item.metadata["refinement_error"] = str(exc)
if on_progress:
on_progress()
return
for item in chunk:
result = by_id.get(item.id)
if result is None:
item.metadata["refinement_error"] = "Batch refinement returned no entry for this item_id."
if on_progress:
on_progress()
continue
prompt, response_text = item.get_prompt_and_response()
new_prompt = result.refined_prompt if (mode in ("all", "prompt_only") and result.refined_prompt) else prompt
new_response = result.refined_response if mode in ("all", "response_only") else response_text
item.update_content(new_prompt=new_prompt, new_response=new_response)
item.metadata["refined"] = True
item.metadata["refinement_info"] = {
"reasoning_added": result.reasoning_added,
"formatting_fixed": result.formatting_fixed,
"artifacts_removed": result.artifacts_removed,
"explanation": result.explanation_of_changes,
}
if on_progress:
on_progress()
async def refine_batch_async(
self,
items: List[DatasetItem],
mode: str = "all",
batch_size: int = 20,
max_field_chars: int = 6000,
on_progress: Optional[callable] = None,
) -> List[DatasetItem]:
"""Refine a batch of items, grouping records into structured requests of up to
`batch_size` so the (fairly long) system prompt is sent once per group instead of
once per record -- cuts both token cost and network round-trips versus one call
per item, with no change in refinement quality per record.
"""
chunks = [items[i:i + batch_size] for i in range(0, len(items), batch_size)]
await asyncio.gather(*(
self._refine_chunk(chunk, mode, max_field_chars, on_progress)
for chunk in chunks
))
return items