import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

model_id = "maharshpatelx/SmolGRPO-135M"

device = 0 if torch.cuda.is_available() else -1
print(f"Using device index: {device}")

model = AutoModelForCausalLM.from_pretrained(model_id)
tokenizer = AutoTokenizer.from_pretrained(model_id)

pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    device=device,
)

messages = [
    {
        "role": "user",
        "content": (
            "Summarize this paragraph in about 50 tokens:\n\n"
            "Cats are small domesticated carnivores that live closely with humans..."
        ),
    },
]

outputs = pipe(
    messages,
    max_new_tokens=128,
    do_sample=True,
    temperature=0.5,
    min_p=0.1,
)

print(outputs["generated_text"])

prompt = """
# A long document about the Cat

The cat (Felis catus), also referred to as the domestic cat or house cat, is a small
domesticated carnivorous mammal. It is the only domesticated species of the family Felidae.
Advances in archaeology and genetics have shown that the domestication of the cat occurred
in the Near East around 7500 BC. It is commonly kept as a pet and farm cat, but also ranges
freely as a feral cat avoiding human contact. It is valued by humans for companionship and
its ability to kill vermin. Its retractable claws are adapted to killing small prey species
such as mice and rats. It has a strong, flexible body, quick reflexes, and sharp teeth,
and its night vision and sense of smell are well developed. It is a social species,
but a solitary hunter and a crepuscular predator. Cat communication includes
vocalizations—including meowing, purring, trilling, hissing, growling, and grunting—as
well as body language. It can hear sounds too faint or too high in frequency for human ears,
such as those made by small mammals. It secretes and perceives pheromones.
"""

messages = [
    {"role": "user", "content": prompt},
]

outputs = pipe(
    messages,
    max_new_tokens=256,
    do_sample=True,
    temperature=0.5,
    min_p=0.1,
)

print(outputs)

ideal_length = 50

def reward_len(completions, **kwargs):
    return [-abs(ideal_length - len(completion)) for completion in completions]

from peft import LoraConfig, get_peft_model

lora_config = LoraConfig(
    task_type="CAUSAL_LM",
    r=16,
    lora_alpha=32,
    target_modules="all-linear",
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()

from trl import GRPOConfig

training_args = GRPOConfig(
    output_dir="GRPO",
    learning_rate=2e-5,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=1,
    max_prompt_length=256,
    max_completion_length=64,
    num_generations=2,
    optim="adamw_torch",
    num_train_epochs=1,
    bf16=True,                  # when supported on GPU
    report_to=["wandb"],
    remove_unused_columns=False,
    logging_steps=1,
)

merged_model = trainer.model.merge_and_unload()
merged_model.push_to_hub("SmolGRPO-135M", private=False, tags=["GRPO", "Reasoning-Course"])