File size: 1,533 Bytes
4001ec8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
from happytransformer import HappyTextToText, TTSettings
import difflib
import torch
grammar_model = HappyTextToText("T5", "vennify/t5-base-grammar-correction")
beam_settings = TTSettings(num_beams=5, min_length=1, max_length=20)
class GrammarChecker:
def __init__(self):
self.grammar_check = grammar_model
def correct_grammar(self, text):
# Generate corrected text from the grammar model
matches = self.grammar_check.generate_text(text, args=beam_settings)
# Extract original and corrected text
original_text = text
corrected_text = matches.text if matches and matches.text else text
# Calculate the differences between original and corrected text
differences = list(difflib.ndiff(original_text.split(), corrected_text.split()))
# Extract corrected words
corrected_words = [word[2:] for word in differences if word.startswith('+ ')]
# Calculate the total count of found mistakes
foundmistakes_count = len(corrected_words)
return print(f"Corrected: {corrected_text}\nWords Corrected: {corrected_words}\nCount:{foundmistakes_count}")
# def save_model(self, path):
# # Save the model using torch.save
# torch.save(self.grammar_check.model, path)
if __name__ == "__main__":
obj = GrammarChecker()
message = "they is going to spent time together."
# corrected_text, corrected_words, count = obj.correct_grammar(message)
obj.correct_grammar(message)
grammar_model.save("F:\\nlp") |