grammar.py

from gramformer import Gramformer, ERROR_EXPLANATIONS


def analyse_grammar_gf(transcript: str,
                       gf: Gramformer = None):
    if gf is None:
        import torch
        gf = Gramformer(models=1, use_gpu=torch.cuda.is_available())

    corrected_sentences = gf.correct(transcript, max_candidates=1)

    corrections = []
    for corrected_sentence in corrected_sentences:
        edits = gf.get_edits(transcript, corrected_sentence)
        for e in edits:
            if e[0] in ERROR_EXPLANATIONS:
                correction = {
                    "issue_type": "Grammar",
                    # R, M, U
                    "category": e[0][0] if e[0][0] in ["R", "M", "U"] else "R",
                    "matched_text": e[1],
                    "offset": e[2],
                    "explaination": ERROR_EXPLANATIONS.get(e[0], ERROR_EXPLANATIONS["R:OTHER"]),
                    "suggestion": e[4]
                }
                corrections.append(correction)
    return corrections