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)