from transformers import T5Tokenizer, T5ForConditionalGeneration

# Load model and tokenizer
tokenizer = T5Tokenizer.from_pretrained("huytd189/grammar-t5-small")
model = T5ForConditionalGeneration.from_pretrained("huytd189/grammar-t5-small")

input_texts = [
    "Rewrite and fix: 'I will drank two bottle'",
    "Rewrite and fix: 'She go to the store'",
    "Rewrite and fix: 'He is more taller than me'",
    "Rewrite and fix: 'The weather in 1816 Europe was abnormally wet, keeping many inhabitants indoors that summer. From April until September of that year, \"it rained in Switzerland on 130 out of the 183 days from April to September\" (Phillips, 2006). Unlike today, one could not simply turn on a television or swipe around the Internet, looking at posts and videos in order to entertain oneself. Instead, it was much more common for the educated people of the day to spend time reading, discussing well-known authors and artists of the day, playing at cards and walking in their gardens and walking paths.'",
    "Paraphrase to make it coherent: I've expierienced wet and reiny days. But I survive back then",
    "Correct the grammar: 'She don't like appless.'"
]

for input_text in input_texts:
    # Tokenize input
    input_ids = tokenizer(input_text, return_tensors="pt").input_ids

    # Generate with better parameters
    outputs = model.generate(
        input_ids,
        max_length=100,
        num_beams=4,
        early_stopping=True
    )

    # Decode output
    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
    print(f"Input: {input_text}")
    print(f"Output: {result}\n")