upb-nlp/gec-ro-comments
Viewer • Updated • 3.97k • 62 • 1

import torch
from transformers import T5Tokenizer, T5ForConditionalGeneration
# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Load model and tokenizer
model_name = "ionut-visan/Flan-T5-Large_Grammar_Ro"
tokenizer = T5Tokenizer.from_pretrained(model_name)
model = T5ForConditionalGeneration.from_pretrained(model_name).to(device)
model.eval()
# Function to correct grammar
def correct_sentence(sentence):
input_text = "grammar: " + sentence
inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
with torch.no_grad():
outputs = model.generate(inputs, max_length=128, num_beams=4, early_stopping=True)
return tokenizer.decode(outputs[0], skip_special_tokens=True)
# Interactive loop
print("Romanian Grammar Corrector (type 'exit' to quit)")
while True:
user_input = input("\nEnter a sentence to correct: ")
if user_input.lower() == "exit":
print("Exiting. 👋")
break
corrected = correct_sentence(user_input)
print("Corrected:", corrected)
Base model
google/flan-t5-large