tinyquery/concrete.py

"""Directly paraphrase concrete training scenarios, keeping the original reference action."""
import argparse
import concurrent.futures
import json
import random
import re
import time
import urllib.request
from pathlib import Path
from tinyquery.data import compact,serialize
from tinyquery.recipes import LANGUAGES


def paraphrase(row,base,seed):
    task={'schema':row['context']['schema'],'backend':row['backend'],'reference_request':row['question'],
          'reference_action':row['target']}
    prompt=('Write one natural user request in each of four styles that has exactly the reference meaning. '
            'Use conversational phrasing, not instructions about writing SQL. Keep every numeric bound, '
            'literal value, selected field, sorting and limit unchanged. Do not add a new condition. '
            'You may use ordinary English/Hindi words for obvious schema fields, but preserve database '
            'literal values exactly. Output JSON keys en, noisy_en, hi, hinglish with string values. '
            'hi must be Devanagari Hindi, hinglish Romanized Hindi, noisy_en imperfect English. Task: '+compact(task))
    payload={'model':'Qwen/Qwen3.8-27B-FP8','messages':[{'role':'user','content':prompt}],
             'max_tokens':500,'temperature':.7,'seed':seed,'response_format':{'type':'json_object'},
             'chat_template_kwargs':{'enable_thinking':False}}
    request=urllib.request.Request(base+'/chat/completions',data=compact(payload).encode(),headers={'Content-Type':'application/json'})
    with urllib.request.urlopen(request,timeout=180) as response: result=json.load(response)
    parsed=json.loads(result['choices'][0]['message']['content'])
    accepted=[]
    for lang in LANGUAGES:
        question=parsed.get(lang)
        if not isinstance(question,str) or not 5<len(question)<900 or '{' in question: continue
        if lang=='hi' and not re.search('[\u0900-\u097f]',question): continue
        if re.search('[\u0600-\u06ff]',question): continue
        # Reject changed numeric bounds and quoted database literals in constrained SQL recipes.
        source_numbers=set(re.findall(r'(?<![a-zA-Z_])\d+(?:\.\d+)?',row['question']))
        output_numbers=set(re.findall(r'(?<![a-zA-Z_])\d+(?:\.\d+)?',question))
        if source_numbers!=output_numbers: continue
        required={'eq':['value'],'project_eq':['value'],'and':['value'],'or':['value','other'],
                  'count_eq':['value'],'contains':['value'],'join_filter':['value']}.get(row['operation'],[])
        if any(str(row['slots'][k]).casefold() not in question.casefold() for k in required): continue
        new=dict(row); new['id']=row['scenario_id']+'_concrete_'+lang; new['language']=lang
        new['question']=question; new['prompt']=serialize(new['context'],question)
        new['provenance']='Qwen3.8-27B-FP8 direct concrete paraphrase; numeric/literal checks; semantic audit separate'
        accepted.append(new)
    return {'scenario_id':row['scenario_id'],'request':payload,'raw_response':result,'accepted':accepted}


def main():
    p=argparse.ArgumentParser(); p.add_argument('--data',required=True); p.add_argument('--out',required=True)
    p.add_argument('--count',type=int,default=2500); p.add_argument('--workers',type=int,default=32)
    p.add_argument('--base',default='http://127.0.0.1:18000/v1'); args=p.parse_args()
    from tinyquery.data import SQL_OPS
    seen=set(); candidates=[]
    with open(args.data) as stream:
        for line in stream:
            r=json.loads(line)
            if r['scenario_id'] not in seen and r['operation'] in SQL_OPS:
                seen.add(r['scenario_id']); candidates.append(r)
    random.Random(887).shuffle(candidates); candidates=candidates[:args.count]
    out=Path(args.out); out.parent.mkdir(parents=True,exist_ok=True)
    done=set()
    if out.exists():
        done={json.loads(line)['scenario_id'] for line in out.read_text().splitlines()}
    start=time.time(); total=0; completed=0
    with out.open('a') as stream, concurrent.futures.ThreadPoolExecutor(args.workers) as pool:
        jobs={pool.submit(paraphrase,r,args.base,2100+i):r['scenario_id'] for i,r in enumerate(candidates) if r['scenario_id'] not in done}
        for future in concurrent.futures.as_completed(jobs):
            completed+=1
            try:
                result=future.result(); stream.write(compact(result)+'\n'); stream.flush()
                total+=len(result['accepted'])
            except Exception as exc:
                print(compact({'error':str(exc),'scenario_id':jobs[future]}),flush=True)
            if completed%25==0: print(compact({'completed':completed,'accepted':total,'seconds':time.time()-start}),flush=True)


if __name__=='__main__': main()