from transformers import AutoModelForCausalLM, AutoTokenizer

# טעינת המודל וה-tokenizer
model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")

while True:
    # קבלת פרומפט מהמשתמש
    prompt = input("Enter your prompt (or type 'exit' to quit): ")
    
    if prompt.lower() == "exit":
        print("Exiting the chatbot.")
        break

    # יצירת input_ids ו-attention_mask מהפרומפט שהוזן
    inputs = tokenizer(prompt, return_tensors="pt")
    input_ids = inputs.input_ids
    attention_mask = inputs.attention_mask  # הוספת attention_mask

    # יצירת טקסט בעזרת המודל
    gen_tokens = model.generate(
        input_ids,
        attention_mask=attention_mask,  # הגדרת attention_mask
        do_sample=True,
        temperature=0.9,
        max_length=100,
        pad_token_id=tokenizer.eos_token_id  # הגדרת pad_token_id ל-eos_token_id
    )

    # פענוח התגובה מהמודל
    gen_text = tokenizer.batch_decode(gen_tokens, skip_special_tokens=True)[0]
    print("Response:", gen_text)