import sys
from transformers import AutoTokenizer, AutoModelForCausalLM

def generate_text(prompt, model_name="models/iraqi_proverb_model", max_length=50):
    # Load the trained model and tokenizer
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name)

    # Encode the prompt and generate text
    inputs = tokenizer.encode(prompt, return_tensors="pt")
    outputs = model.generate(
        inputs,
        max_length=max_length,
        num_return_sequences=1,
        do_sample=True,
        top_k=50,
        pad_token_id=tokenizer.eos_token_id  # Set pad_token_id to eos_token_id to avoid warnings
    )

    # Decode the output and print the generated text
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return generated_text

# Main script
if __name__ == "__main__":
    # Get prompt from command-line arguments, default to an example prompt if none provided
    prompt = sys.argv[1] if len(sys.argv) > 1 else "في العراق، يقولون"
    print("Generated text:", generate_text(prompt))