| import os,json,time,random,gc,hashlib |
| from pathlib import Path |
| import torch |
| from datasets import load_dataset |
| from transformers import AutoTokenizer,AutoModelForCausalLM,Trainer,TrainingArguments,set_seed |
| from peft import LoraConfig,get_peft_model |
| from huggingface_hub import HfApi,hf_hub_download |
| from sacrebleu.metrics import CHRF,BLEU |
| BASE='utter-project/EuroLLM-1.7B';REPO='SEE7/AstheralTranslate';DATA='Helsinki-NLP/opus-100' |
| api=HfApi();rev=api.model_info(BASE).sha;datarev=api.dataset_info(DATA).sha |
| out=Path('/tmp/translate');out.mkdir(exist_ok=True) |
| set_seed(42);rng=random.Random(42);start=time.monotonic() |
| langs={'de':'German','fr':'French','es':'Spanish','it':'Italian','pt':'Portuguese'} |
| tok=AutoTokenizer.from_pretrained(BASE,revision=rev);tok.pad_token=tok.eos_token;tok.padding_side='left' |
| def prompt(text,lang):return 'Translate from English to '+langs[lang]+'. Return only the translation.\nEnglish: '+text.strip()+'\n'+langs[lang]+':' |
| train=[];tests=[];seen=set() |
| for lang in langs: |
| conf='-'.join(sorted(['en',lang])) |
| ds=load_dataset(DATA,conf,split='train',revision=datarev) |
| ds=ds.shuffle(seed=42) |
| count=0 |
| for row in ds: |
| x=row['translation'];a=x['en'].strip();b=x[lang].strip() |
| if not (8<=len(a)<=600 and 4<=len(b)<=700):continue |
| key=(lang,a,b) |
| if key in seen:continue |
| pre=tok(prompt(a,lang),add_special_tokens=True)['input_ids'] |
| answer=tok(' '+b+tok.eos_token,add_special_tokens=False)['input_ids'] |
| if len(pre)+len(answer)>256:continue |
| seen.add(key);train.append({'input_ids':pre+answer,'labels':[-100]*len(pre)+answer,'source':a,'target':b,'language':lang}) |
| count+=1 |
| if count==2000:break |
| print('DATA',lang,count,flush=True) |
| test=load_dataset(DATA,conf,split='test',revision=datarev) |
| for row in test: |
| x=row['translation'];a=x['en'].strip();b=x[lang].strip() |
| if 8<=len(a)<=300 and 4<=len(b)<=400 and (lang,a,b) not in seen: |
| tests.append({'source':a,'reference':b,'language':lang}) |
| if sum(t['language']==lang for t in tests)==25:break |
| assert len(tests)==125 |
| rng.shuffle(train) |
| (out/'training_manifest.json').write_text(json.dumps({'base':BASE,'revision':rev,'dataset':DATA,'dataset_revision':datarev,'examples':len(train),'per_language':2000,'seed':42,'sha256':hashlib.sha256(json.dumps(train,sort_keys=True).encode()).hexdigest()},indent=2)) |
| model=AutoModelForCausalLM.from_pretrained(BASE,revision=rev,torch_dtype=torch.bfloat16,device_map={'':'cuda'},attn_implementation='sdpa').eval() |
| report={'base':BASE,'base_revision':rev,'dataset_revision':datarev,'parameters':sum(p.numel() for p in model.parameters()),'languages':langs,'hardware':torch.cuda.get_device_name(),'training_examples':len(train)} |
| @torch.inference_mode() |
| def evaluate(label): |
| rows=[] |
| for i,t in enumerate(tests): |
| x=tok(prompt(t['source'],t['language']),return_tensors='pt').to('cuda') |
| y=model.generate(**x,max_new_tokens=128,do_sample=False,pad_token_id=tok.pad_token_id,eos_token_id=tok.eos_token_id) |
| raw=tok.decode(y[0,x['input_ids'].shape[1]:],skip_special_tokens=True) |
| pred=raw.strip().split('\n')[0].strip() |
| rows.append({**t,'prediction':pred,'raw':raw,'generated_tokens':int(y.shape[1]-x['input_ids'].shape[1])}) |
| if i%10==0:print('EVAL',label,i,flush=True) |
| scores={} |
| for lang in langs: |
| r=[r for r in rows if r['language']==lang];hyp=[r['prediction'] for r in r];refs=[[r['reference'] for r in r]] |
| scores[lang]={'chrf':CHRF().corpus_score(hyp,refs).score,'bleu':BLEU().corpus_score(hyp,refs).score,'n':len(r)} |
| (out/(label+'_outputs.json')).write_text(json.dumps(rows,ensure_ascii=False,indent=2)) |
| print('SCORES',label,json.dumps(scores),flush=True) |
| return scores |
| report['base_scores']=evaluate('base') |
| model=get_peft_model(model,LoraConfig(r=32,lora_alpha=64,lora_dropout=0.05,target_modules=['q_proj','k_proj','v_proj','o_proj','gate_proj','up_proj','down_proj'],task_type='CAUSAL_LM')) |
| report['trainable_parameters']=sum(p.numel() for p in model.parameters() if p.requires_grad) |
| model.config.use_cache=False |
| def collate(rows): |
| length=max(len(r['input_ids']) for r in rows) |
| return {'input_ids':torch.tensor([[tok.pad_token_id]*(length-len(r['input_ids']))+r['input_ids'] for r in rows]),'labels':torch.tensor([[-100]*(length-len(r['labels']))+r['labels'] for r in rows]),'attention_mask':torch.tensor([[0]*(length-len(r['input_ids']))+[1]*len(r['input_ids']) for r in rows])} |
| args=TrainingArguments(output_dir='/tmp/checkpoints',max_steps=500,per_device_train_batch_size=4,gradient_accumulation_steps=4,learning_rate=1e-4,warmup_ratio=0.05,bf16=True,gradient_checkpointing=True,gradient_checkpointing_kwargs={'use_reentrant':False},save_strategy='no',logging_steps=20,report_to='none',remove_unused_columns=False,seed=42) |
| trainer=Trainer(model=model,args=args,train_dataset=train,data_collator=collate) |
| result=trainer.train();report['training']=result.metrics;report['optimizer_steps']=trainer.state.global_step |
| model.save_pretrained(out/'adapter');model.gradient_checkpointing_disable();model.config.use_cache=True;model.eval() |
| model=model.merge_and_unload(safe_merge=True);model.config.astheral_version='translate-1' |
| model.save_pretrained(out/'model',safe_serialization=True,max_shard_size='4GB');tok.save_pretrained(out/'model') |
| del trainer,model;gc.collect();torch.cuda.empty_cache() |
| model=AutoModelForCausalLM.from_pretrained(out/'model',torch_dtype=torch.bfloat16,device_map={'':'cuda'},attn_implementation='sdpa').eval() |
| report['finetuned_scores']=evaluate('finetuned') |
| report['mean_chrf_base']=sum(v['chrf'] for v in report['base_scores'].values())/5 |
| report['mean_chrf_finetuned']=sum(v['chrf'] for v in report['finetuned_scores'].values())/5 |
| report['improved']=report['mean_chrf_finetuned']>report['mean_chrf_base'] |
| report['elapsed_seconds']=time.monotonic()-start |
| (out/'evaluation.json').write_text(json.dumps(report,indent=2)) |
| (out/'train.py').write_text(Path(__file__).read_text()) |
| api.upload_folder(repo_id=REPO,folder_path=str(out/'model'),commit_message='Publish actual merged AstheralTranslate weights after reload and translation evaluation') |
| api.upload_folder(repo_id=REPO,folder_path=str(out/'adapter'),path_in_repo='adapter',commit_message='Publish trained rank32 translation adapter') |
| api.upload_folder(repo_id=REPO,folder_path=str(out),allow_patterns=['*.json','*.py'],commit_message='Publish translation training provenance and before/after test outputs') |
| import urllib.request |
| license_path=str(out/'LICENSE') |
| Path(license_path).write_bytes(urllib.request.urlopen('https://www.apache.org/licenses/LICENSE-2.0.txt',timeout=20).read()) |
| api.upload_file(repo_id=REPO,path_or_fileobj=license_path,path_in_repo='LICENSE') |
| print('TRANSLATE_COMPLETE',json.dumps(report),flush=True) |
|
|
