gramformer_cli.py

# A script for using https://github.com/PrithivirajDamodaran/Gramformer
# Write by GPT-3.5🧙, scillidan🤡
# see https://shareg.pt/9aQy8UZ

import sys
import warnings
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from gramformer import Gramformer

def set_seed(seed):
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

def main():
    if len(sys.argv) != 2:
        print("Usage: python cli.py INPUT")
        return

    set_seed(1212)

    warnings.filterwarnings("ignore", category=FutureWarning)

    use_gpu = torch.cuda.is_available()
    device = torch.device("cuda" if use_gpu else "cpu")
    gf = Gramformer(models=1, use_gpu=use_gpu)  # 1=corrector, 2=detector

    influent_sentence = sys.argv[1]
    corrected_sentences = gf.correct(influent_sentence, max_candidates=1)
    print("[Input] ", influent_sentence)
    for corrected_sentence in corrected_sentences:
        print("[Correction] ", corrected_sentence)

    print("Device:", device)

if __name__ == "__main__":
    main()