Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
3device = "cuda" if torch.cuda.is_available() else "cpu"
4MODEL_NAME = "DeepMount00/OCR_corrector"
5
6model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME).eval()
7tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
8model.to(device)
9my_text = "In un'epca lontnaa, un re goernava le sue tere con saggez2a e giustiia. Sotot il suo regno, il rgeno prosperava e la getne era flice. Ma un gionro, un drgoa feroce attcò il regno, semniando ditruzione e paurra tra i suoi abtanti."
10inputs = tokenizer(my_text, return_tensors="pt").to(device)
11outputs = model.generate(input_ids=inputs['input_ids'],
12 attention_mask=inputs['attention_mask'],
13 num_beams=2, max_length=1050, top_k=10)
14clean_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
15print(clean_text)