"""Standalone inference example for angeldust007/tr-factbench-gemma-atomizer."""
import json
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel

BASE_MODEL = "google/gemma-4-E2B-it"
ADAPTER_REPO = "angeldust007/tr-factbench-gemma-atomizer"

SYSTEM_PROMPT = (
    "Türkçe bir iddiayı bağımsız doğrulanabilir atomik önermelere ayır. "
    "Her atom tek başına anlaşılır olmalı. Ortak özne, nesne veya tamlayıcı "
    "ikinci bir önermede düşürülmüşse, onu iddianın kendi içindeki bilgiden yeniden kur. "
    "Özneleri veya baş ögeleri yeniden kurarken kelime eksiltme, kısaltma veya kırpma yapma; "
    "iddiadaki isim ve sıfat tamlamalarını tam olarak koru. "
    "Her atom dilbilgisel olarak eksiksiz ve öznesi açık bir cümle olmalıdır. "
    "Yeni bilgi ekleme. Olumsuzluk, koşul, modalite, sayı, zaman, karşılaştırma ve "
    "kapsam ifadelerini koru. Koşullu tek bir önerme sırf birden fazla fiil içeriyor diye "
    "bölünmemelidir. İddia zaten atomikse değiştirmeden tek atom döndür. "
    "BİR cümlenin içindeki yan cümleler veya sıfat tamlamaları ayrı atom YAPILMAZ; "
    "yalnızca 've', 'ayrıca', 'öte yandan', 'bununla birlikte', 'ancak', 'fakat', "
    "noktalı virgül veya açık nokta gibi koordinasyon sınırlarında böl. "
    'Yalnızca {"atoms": ["..."]} biçiminde geçerli JSON üret. /no_think'
)


def main():
    compute_dtype = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float16
    quant_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_use_double_quant=True,
        bnb_4bit_compute_dtype=compute_dtype,
    )

    print(f"Loading tokenizer & adapter from {ADAPTER_REPO}...")
    tokenizer = AutoTokenizer.from_pretrained(ADAPTER_REPO)
    base_model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        quantization_config=quant_config,
        device_map="auto",
        torch_dtype=compute_dtype,
    )
    model = PeftModel.from_pretrained(base_model, ADAPTER_REPO)
    model.eval()

    sample_claims = [
        "Alışverişte müşteri, işyerinin oluşturduğu karekodu mobil uygulamasından okutarak FAST ödemesi başlatabilir.",
        "Kasko sigortasında Dar Kasko teminat gruplarının yalnız bir bölümünü, Kasko ise tamamını kapsar; zorunlu trafik sigortası üçüncü kişilere verilen maddi ve bedensel zararları güvence altına alır.",
    ]

    for claim in sample_claims:
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": json.dumps({"claim": claim}, ensure_ascii=False)},
        ]
        prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        with torch.no_grad():
            output_ids = model.generate(**inputs, max_new_tokens=256, do_sample=False)
        generated = tokenizer.decode(output_ids[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
        print(f"\nCLAIM: {claim}")
        print(f"OUTPUT: {generated.strip()}")


if __name__ == "__main__":
    main()