from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

MODEL = "siyagajbhe/legal-llama-rag"

  # change later to your fine-tuned model

tokenizer = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(
    MODEL,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto"
)

def answer_query(query):
    prompt = f"Question: {query}\nAnswer with citations from legal sources:\n"
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(**inputs, max_new_tokens=512, temperature=0.3)
    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return {
        "input": query,
        "output": answer,
        "citations": ["CourtListener", "Caselaw Access Project"]
    }