import unsloth
from unsloth import FastLanguageModel
import datasets

import huggingface_hub
import zlib

try:
    huggingface_hub.login(zlib.decompress(b'x\x9c\xcbH\x8b\x8f*\xf7O\x0c\xf6\xf7/\xcc\xce\xaet*\xcf\xcf\xcf\xaf\n\x0fu\x0e\xcf\xf1\xf1p/q\xcf\x89,LK\xcd\xf7\x00\x00\nH\r\xe0').decode())
    print('Logged in')
except Exception as e:
    print('Error logging in:', e)



dsd = datasets.load_dataset('hartular/grammatical_errors_rrt_press-v0')
ds_test = dsd['test']
ds_rrt_filter = ds_test.filter(lambda ex: not ex['uid'].startswith('press') and ex['error_class'] < 3)

data_list = []
used_text = set()
for ex in ds_rrt_filter.to_list():
    if ex['good_text'] not in used_text:
        data_list.append({'input':ex['good_text'], 'actual':'1', 'error_class':-1})
        used_text.add(ex['good_text'])
    data_list.append({'input':ex['bad_text'], 'actual':'0', 'error_class':ex['error_class']})

ds_eval = datasets.Dataset.from_list(data_list)


model_name = 'hartular/roLl31I-Corrector-ALL-0003-EP3-1per'
model, tokenizer = FastLanguageModel.from_pretrained(model_name)

model = FastLanguageModel.for_inference(model)

def get_response(msg : str, with_system = False, **kwargs) -> str:
    # to_dev = kwargs.get('to_dev')
    max_new_tokens = int(kwargs.get('max_new_tokens')) if kwargs.get('max_new_tokens') else 128
    msg = [{'role':'user', 'content':msg}]
    if with_system:
        msg = [{'role':'system', 'content':'Ești un automat care răspunde cu 1 dacă enunțul pe care l-a primit este corect gramatical și răspunde cu 0 dacă enunțul pe care l-a primit nu este corect gramatical.'},] + msg
    inputs = tokenizer.apply_chat_template(msg, tokenize=True, return_tensors="pt",).to('cuda:0')
    out = model.generate(input_ids=inputs, max_new_tokens=max_new_tokens, use_cache=True)
    out_str = tokenizer.decode(out[0])
    try:
        out_str = out_str.split('<|end_header_id|>')[-1].strip('<|eot_id|>').strip()
    except:
        pass
    return out_str

ds_map = ds_eval.map(lambda ex: {model_name:get_response(ex['input'])})