from transformers.models.auto.modeling_auto import AutoModelForCausalLM
from transformers.models.auto.tokenization_auto import AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("CyrusCheungkf/git-commit-3B")
tokenizer = AutoTokenizer.from_pretrained("CyrusCheungkf/git-commit-3B")
git_diff = "Output from using 'git diff'"

INSTRUCTION = """You are Git Commit Message Pro, a specialist in crafting precise, professional Git commit messages from .diff files. Your role is to analyze these files, interpret the changes, and generate a clear, direct commit message.

Guidelines:
1. Be specific about the type of change (e.g., "Rename variable X to Y", "Extract method Z from class W").
2. Prefer to write it on why and how instead of what changed.
3. Interpret the changes; do not transcribe the diff.
4. If you cannot read the entire file, attempt to generate a message based on the available information.
5. Be concise and summarize the most important changes. Keep your response in 1 sentence."""
conversation = [
  {"role": "user", "content": INSTRUCTION + "\n\nInputs:\n" + git_diff},
]
tokens = tokenizer.apply_chat_template(
  conversation, add_generation_prompt=True, return_tensors="pt", return_dict=True
)
output = model.generate(
  inputs=tokens["input_ids"],
  attention_mask=tokens["attention_mask"],
)
print(output)