# -----------------------------------------------------------
# 1. TEXT GENERATION USING HUGGINGFACE PIPELINE
# -----------------------------------------------------------
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
from evaluate import load

# Pipeline for easy text generation
pipe = pipeline(
    "text-generation",
    model="varma007ut/Indian_Legal_Assitant",
    device="cpu",
    max_length=200
)

prompt1 = "Summarize the key points of the Indian Contract Act, 1872:"
result = pipe(prompt1)

print("\n=== Pipeline Output ===")
print(result[0]["generated_text"])


# -----------------------------------------------------------
# 2. TEXT GENERATION USING model.generate()
# -----------------------------------------------------------

tokenizer = AutoTokenizer.from_pretrained("varma007ut/Indian_Legal_Assitant")
model = AutoModelForCausalLM.from_pretrained("varma007ut/Indian_Legal_Assitant")

prompt2 = "What are the fundamental rights in the Indian Constitution?"
inputs = tokenizer(prompt2, return_tensors="pt")

outputs = model.generate(
    **inputs,
    max_length=200,
    do_sample=True,
    top_p=0.95
)

generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)

print("\n=== HF Model.generate Output ===")
print(generated_text)


# -----------------------------------------------------------
# 3. BLEU SCORE EVALUATION (SAFE FOR SPACES)
# -----------------------------------------------------------

bleu = load("bleu")

pred_text = generated_text  # prediction from HF model

# Reference answer (your true dataset reference goes here)
references = ["Fundamental Rights are guaranteed by the Indian Constitution."]

# BLEU expects tokenized text (split into words)
results = bleu.compute(
    predictions=[pred_text.split()],
    references=[[ref.split() for ref in references]]
)

print("\n=== BLEU Score ===")
print(results)