Views
No views yet
1from transformers import AutoTokenizer, AutoModelWithLMHead, AutoModelForCausalLM
2import torch
3if torch.cuda.is_available():
4 device = torch.device("cuda")
5else :
6 device = "cpu"
7
8
9tokenizer = AutoTokenizer.from_pretrained("Ashishkr/grammar_correction")
10model = AutoModelForCausalLM.from_pretrained("Ashishkr/grammar_correction").to(device)
11
12input_query="what be the reason for everyone leave the company"
13query= "<|startoftext|> " + input_query + " ~~~"
14
15
16input_ids = tokenizer.encode(query.lower(), return_tensors='pt').to(device)
17sample_outputs = model.generate(input_ids,
18 do_sample=True,
19 num_beams=1,
20 max_length=128,
21 temperature=0.9,
22 top_p= 0.7,
23 top_k = 5,
24 num_return_sequences=3)
25corrected_sentences = []
26for i in range(len(sample_outputs)):
27 r = tokenizer.decode(sample_outputs[i], skip_special_tokens=True).split('||')[0]
28 r = r.split('~~~')[1]
29 if r not in corrected_sentences:
30 corrected_sentences.append(r)
31
32print(corrected_sentences)