import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
MODEL = "overthelex/ua-legal-citation-grounded-14b"
tok = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.bfloat16, device_map="auto").eval()
SYSTEM = (
"Ти — український юридичний асистент. Відповідай ВИКЛЮЧНО на основі наданих витягів "
"із судових рішень ЄДРСР. Кожне фактичне твердження підкріплюй посиланням у форматі "
"[doc:ID], де ID — це edrsr_doc_id відповідного джерела. Якщо у наданих джерелах немає "
"достатньої підстави — прямо напиши, що наданих джерел недостатньо. Пиши українською."
)
context = [
{"doc_id": "114714815", "text": "...витяг із рішення суду..."},
{"doc_id": "77050963", "text": "...ще один витяг..."},
]
question = "Як суд розподіляє спільне майно подружжя при розлученні?"
ctx = "\n\n".join(f"[doc:{c['doc_id']}] {c['text']}" for c in context)
user = (f"Питання: {question}\n\nДжерела (витяги з рішень ЄДРСР):\n{ctx}\n\n"
"Дай обґрунтовану відповідь українською з посиланнями [doc:ID].")
msgs = [{"role": "system", "content": SYSTEM}, {"role": "user", "content": user}]
ids = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt",
truncation=True, max_length=8192).to(model.device)
out = model.generate(ids, max_new_tokens=512, do_sample=False, pad_token_id=tok.pad_token_id)
answer = tok.decode(out[0, ids.shape[1]:], skip_special_tokens=True)
print(answer)