import json, pandas as pd
from jsonschema import Draft7Validator
from transformers import AutoTokenizer

MODEL_NAME = "Qwen/Qwen2.5-0.5B"
MAX_TOKENS = 4096
N_TRAIN, N_TEST = 1000, 200
SEED = 42

SYSTEM = (
    "You are an information extraction engine. Given a source text and a JSON Schema, "
    "extract the relevant information and return a single JSON object that strictly conforms "
    "to the schema. Output only the JSON, with no extra commentary, markdown, or code fences."
)
INSTRUCTION = (
    "Extract the information from the text below into a JSON object that strictly conforms "
    "to the given JSON Schema. Use only information present in the text. Return only the JSON object."
)

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

def as_obj(x):
    return x if isinstance(x, (dict, list)) else json.loads(x)

def to_chat(text, schema, target):
    user = (f"{INSTRUCTION}\n\n"
            f"### JSON Schema\n{json.dumps(schema, ensure_ascii=False)}\n\n"
            f"### Text\n{str(text).strip()}")
    return {"messages": [
        {"role": "system",    "content": SYSTEM},
        {"role": "user",      "content": user},
        {"role": "assistant", "content": json.dumps(target, ensure_ascii=False)},
    ]}

def fits(rec):
    return len(tokenizer.apply_chat_template(rec["messages"], tokenize=True)) <= MAX_TOKENS

df = pd.read_parquet("hf://datasets/scrapegraphai/scrapegraphai-100k/data/train.parquet")
df = df[df["response_is_valid"]].sample(frac=1, random_state=SEED)

train = []
for _, r in df.iterrows():
    if len(train) >= N_TRAIN:
        break
    try:
        schema, target = as_obj(r["schema"]), as_obj(r["response"])
        if not str(r["content"]).strip():
            continue
        Draft7Validator(schema).validate(target)
        rec = to_chat(r["content"], schema, target)
        if fits(rec):
            train.append(rec)
    except Exception:
        continue

dt = pd.read_parquet("hf://datasets/paraloq/json_data_extraction/data.parquet").sample(frac=1, random_state=SEED)

test = []
for _, r in dt.iterrows():
    if len(test) >= N_TEST:
        break
    try:
        rec = to_chat(r["text"], as_obj(r["schema"]), as_obj(r["item"]))
        if fits(rec):
            test.append(rec)
    except Exception:
        continue

import json, random
from datasets import load_dataset, get_dataset_config_names

LETTERS = ["A", "B", "C", "D"]
N_PER_SUBJECT = 5
SEED = 42

subjects = [c for c in get_dataset_config_names("cais/mmlu") if c not in {"all", "auxiliary_train"}]

def to_prompt(question, choices):
    opts = "\n".join(f"{L}. {c}" for L, c in zip(LETTERS, choices))
    return ("Answer the following multiple choice question. "
            "Respond with only the letter (A, B, C, or D) of the correct answer.\n\n"
            f"Question: {question}\n{opts}\nAnswer:")

mmlu = []
for subj in subjects:
    ds = load_dataset("cais/mmlu", subj, split="test")
    idx = list(range(len(ds)))
    random.Random(SEED).shuffle(idx)
    for i in idx[:N_PER_SUBJECT]:
        ex = ds[i]
        mmlu.append({
            "subject": subj,
            "answer_letter": LETTERS[ex["answer"]],
            "messages": [{"role": "user", "content": to_prompt(ex["question"], ex["choices"])}],
        })

def save_jsonl(path, rows):
    with open(path, "w", encoding="utf-8") as f:
        for x in rows:
            f.write(json.dumps(x, ensure_ascii=False) + "\n")

save_jsonl("train.jsonl", train)
save_jsonl("test.jsonl",  test)
save_jsonl("mmlu_probe.jsonl", mmlu)

# !pip install -q -U transformers peft trl datasets accelerate
# !pip install -q jsonschema python-Levenshtein matplotlib

import os, gc, json, random, re
from contextlib import contextmanager
from collections import defaultdict

import numpy as np
import torch
import Levenshtein
import matplotlib.pyplot as plt
from tqdm.auto import tqdm

from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
from peft import LoraConfig, get_peft_model
from jsonschema import Draft7Validator

SEED = 42
random.seed(SEED); np.random.seed(SEED)
torch.manual_seed(SEED); torch.cuda.manual_seed_all(SEED)

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

MODEL_NAME = "Qwen/Qwen2.5-0.5B"
MAX_LEN = 4096


tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"


model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16, device_map={"": 0})
model.config.use_cache = False
model.config.pad_token_id = tokenizer.pad_token_id


print(f"{sum(p.numel() for p in model.parameters())/1e6:.0f}M параметров")

raw = load_dataset("json", data_files={"train": "train.jsonl", "test": "test.jsonl"})
split = raw["train"].train_test_split(test_size=0.1, seed=SEED)


train_ds, val_ds, test_ds = split["train"], split["test"], raw["test"]
mmlu = [json.loads(l) for l in open("mmlu_probe.jsonl", encoding="utf-8") if l.strip()]


print("train:", len(train_ds), "| val:", len(val_ds), "| test:", len(test_ds), "| mmlu:", len(mmlu))
# train: 900 | val: 100 | test: 200 | mmlu: 285 

def extract_schema(messages):
    user = next(m["content"] for m in messages if m["role"] == "user")
    return json.loads(user.split("### JSON Schema", 1)[1].split("### Text", 1)[0].strip())


def gold(messages):
    return next(m["content"] for m in messages if m["role"] == "assistant")


def _first_span(t):
    start = next((i for i, c in enumerate(t) if c in "{["), None)
    if start is None: return None
    stack, pairs = [], {"}": "{", "]": "["}
    for j in range(start, len(t)):
        c = t[j]
        if c in "{[": stack.append(c)
        elif c in "}]":
            if not stack or stack.pop() != pairs[c]: return None
            if not stack: return t[start:j+1]
    return None


def parse_json(text):
    t = re.sub(r"^```(?:json)?|```$", "", text.strip()).strip()
    for cand in (t, _first_span(t)):
        if cand is None: continue
        try:
            obj = json.loads(cand)
            return obj, json.dumps(obj, sort_keys=True, ensure_ascii=False)
        except Exception:
            pass
    return None, ""


def score(preds, examples):
    n = len(preds); valid = 0; lev = []
    for p, ex in zip(preds, examples):
        _, g_canon = parse_json(gold(ex["messages"]))
        p_obj, p_canon = parse_json(p)
        if p_obj is not None:
            try:
                Draft7Validator(extract_schema(ex["messages"])).validate(p_obj); valid += 1
            except Exception:
                pass
        lev.append(Levenshtein.ratio(p_canon if p_obj is not None else p.strip(), g_canon))
    return {"schema_valid": valid/n, "levenshtein": sum(lev)/n}

LETTERS = ["A", "B", "C", "D"]

@contextmanager
def left_padding(tok):
    old = tok.padding_side; tok.padding_side = "left"
    try: yield
    finally: tok.padding_side = old

@torch.no_grad()
def generate(examples, batch_size=24, max_new_tokens=1024):
    model.eval()
    cache = model.config.use_cache;
    model.config.use_cache = True
    out = []
    try:
        with left_padding(tokenizer):
            for i in tqdm(range(0, len(examples), batch_size), desc="generate"):
                batch = examples[i:i+batch_size]
                prompts = [tokenizer.apply_chat_template(
                    [m for m in ex["messages"] if m["role"] != "assistant"],
                    tokenize=False, add_generation_prompt=True) for ex in batch]
                enc = tokenizer(prompts, return_tensors="pt", padding=True,
                                truncation=True, max_length=MAX_LEN).to(model.device)
                gen = model.generate(**enc, max_new_tokens=max_new_tokens, do_sample=False,
                                     use_cache=True, pad_token_id=tokenizer.pad_token_id)
                out += tokenizer.batch_decode(gen[:, enc["input_ids"].shape[1]:], skip_special_tokens=True)
    finally:
        model.config.use_cache = cache
    return out

@torch.no_grad()
def eval_mmlu(data, batch_size=4):
    model.eval()
    letter_ids = [tokenizer(" " + L, add_special_tokens=False).input_ids[0] for L in LETTERS]
    corr = total = 0
    with left_padding(tokenizer):
        for i in tqdm(range(0, len(data), batch_size), desc="mmlu"):
            batch = data[i:i+batch_size]
            enc = tokenizer([ex["messages"][0]["content"] for ex in batch], return_tensors="pt",
                            padding=True, truncation=True, max_length=MAX_LEN).to(model.device)
            logits = model(**enc).logits[:, -1, :][:, letter_ids]
            for ex, pi in zip(batch, logits.argmax(-1).tolist()):
                corr += int(LETTERS[pi] == ex["answer_letter"]); total += 1
    return corr / total

test_examples = [test_ds[i] for i in range(len(test_ds))]


base_json = score(generate(test_examples), test_examples)
base_mmlu = eval_mmlu(mmlu)


baseline = {**base_json, "mmlu": base_mmlu}
print(baseline)
# {'schema_valid': 0.265,'levenshtein': 0.4112473125422889, 'mmlu': 0.45964912280701753} 

# ! pip uninstall torchao
model.gradient_checkpointing_enable()
model.enable_input_require_grads()

lora = LoraConfig(
    r=16, lora_alpha=32, lora_dropout=0.05, bias="none", task_type="CAUSAL_LM",
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
)
model = get_peft_model(model, lora)

def tokenize_and_mask(example):
    msgs = example["messages"]
    full   = tokenizer.apply_chat_template(msgs, tokenize=False)
    prefix = tokenizer.apply_chat_template(
        [m for m in msgs if m["role"] != "assistant"], tokenize=False, add_generation_prompt=True)
    ids = tokenizer(full, add_special_tokens=False, truncation=True, max_length=MAX_LEN)["input_ids"]
    n = min(len(tokenizer(prefix, add_special_tokens=False)["input_ids"]), len(ids))
    labels = [-100] * n + ids[n:]            # loss только на ответе ассистента
    return {"input_ids": ids, "labels": labels}

train_tok = train_ds.map(tokenize_and_mask, remove_columns=train_ds.column_names)
val_tok   = val_ds.map(tokenize_and_mask,   remove_columns=val_ds.column_names)

VAL_PROBE = [val_ds[i] for i in range(min(40, len(val_ds)))]
history = [{"step": 0, "schema_valid": baseline["schema_valid"],
            "levenshtein": baseline["levenshtein"], "mmlu": baseline["mmlu"]}]

def preprocess_logits_for_metrics(logits, labels):
    if isinstance(logits, tuple): logits = logits[0]
    return logits.argmax(dim=-1)

def compute_metrics(eval_pred):
    gc_on = model.is_gradient_checkpointing
    if gc_on: model.gradient_checkpointing_disable()
    
    jm = score(generate(VAL_PROBE, batch_size=8, max_new_tokens=512), VAL_PROBE)
    macc = eval_mmlu(mmlu, batch_size=4)
    
    if gc_on: model.gradient_checkpointing_enable()
    
    out = {"schema_valid": jm["schema_valid"], "levenshtein": jm["levenshtein"], "mmlu": macc}
    history.append({"step": trainer.state.global_step, **out})
    return out

class Collator:
    def __call__(self, feats):
        m = max(len(f["input_ids"]) for f in feats); pad = tokenizer.pad_token_id
        ids = [f["input_ids"] + [pad]*(m-len(f["input_ids"])) for f in feats]
        att = [[1]*len(f["input_ids"]) + [0]*(m-len(f["input_ids"])) for f in feats]
        lab = [f["labels"] + [-100]*(m-len(f["labels"])) for f in feats]
        return {"input_ids": torch.tensor(ids), "attention_mask": torch.tensor(att), "labels": torch.tensor(lab)}

args = TrainingArguments(
    output_dir="qwen-lora-json",
    num_train_epochs=1,
    per_device_train_batch_size=1, 
    gradient_accumulation_steps=8, 
    per_device_eval_batch_size=1,
    learning_rate=2e-4, 
    lr_scheduler_type="cosine", 
    warmup_ratio=0.03, 
    fp16=True,
    eval_strategy="steps", 
    eval_steps=25,
    save_strategy="steps", 
    save_steps=25, 
    save_total_limit=2, 
    eval_accumulation_steps=1,
    logging_steps=10,
    load_best_model_at_end=True, 
    metric_for_best_model="eval_loss", 
    greater_is_better=False,
    remove_unused_columns=False, 
    report_to="none",
)

trainer = Trainer(
    model=model, args=args,
    train_dataset=train_tok, eval_dataset=val_tok,
    data_collator=Collator(),
    compute_metrics=compute_metrics,
    preprocess_logits_for_metrics=preprocess_logits_for_metrics,
)