Views
No views yet
python predict.py --input <raw text file> --restore_dir gotutiyan/gector-roberta-base-5k --out <path to output file>1from transformers import AutoTokenizer
2from gector.modeling import GECToR
3from gector.predict import predict, load_verb_dict
4import torch
5
6model_id = 'gotutiyan/gector-roberta-base-5k'
7model = GECToR.from_pretrained(model_id)
8if torch.cuda.is_available():
9 model.cuda()
10tokenizer = AutoTokenizer.from_pretrained(model_id)
11encode, decode = load_verb_dict('data/verb-form-vocab.txt')
12srcs = [
13 'This is a correct sentence.',
14 'This are a wrong sentences'
15]
16corrected = predict(
17 model, tokenizer, srcs,
18 encode, decode,
19 keep_confidence=0.0,
20 min_error_prob=0.0,
21 n_iteration=5,
22 batch_size=2,
23)
24print(corrected)