from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

model_path = "PATH_MODEL_FILES"

tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    dtype=torch.float32
)

SYSTEM_PROMPT = "You are JibayAi. Reply in ONE short sentence."

def generate_reply(user_text):
    prompt = f"{SYSTEM_PROMPT}\nUser: {user_text}\nAssistant:"

    inputs = tokenizer(prompt, return_tensors="pt")

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=25,
            do_sample=False,
            repetition_penalty=1.05,
            pad_token_id=tokenizer.eos_token_id
        )

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

    answer = decoded.split("Assistant:")[-1]
    answer = answer.split("\n")[0].strip()

    return answer


while True:
    user_input = input("You: ")
    if user_input.lower() == "exit":
        break

    reply = generate_reply(user_input)
    print("JibayAi:", reply)
    print("-" * 40)