gramformer_cli_demo.py

# Refer to https://deeplx.owo.network/integration/python.html
# Write by GPT-4o mini 👨‍💻, scillidan 🤡

# Install
# git clone --depth=1 https://github.com/PrithivirajDamodaran/Gramformer
# python -m venv venv
# venv\Scripts\activate.bat
# pip install torch --index-url https://download.pytorch.org/whl/cu121
# pip install -e .
# python -m spacy download en

# How to use
# python this.py "<write>"

import sys
import warnings
from gramformer import Gramformer
import torch

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

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

def main(input_sentence):
    set_seed(1212)

    gf = Gramformer(models=1, use_gpu=torch.cuda.is_available())  # 1=corrector, 2=detector

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

    print("x ", input_sentence)
    if corrected_sentences:
        for corrected_sentence in corrected_sentences:
            print("v ", corrected_sentence)
    else:
        print("No corrections found.")
    print("Device:", "cuda" if torch.cuda.is_available() else "cpu")

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: python cli.py <input>")
        sys.exit(1)
    
    input_sentence = sys.argv[1]
    main(input_sentence)