Update README.md

README.md

@@ -163,25 +163,15 @@ We compare our solution with both open automatic spell checkers and the ChatGPT
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
 tokenizer = AutoTokenizer.from_pretrained("ai-forever/sage-fredt5-distilled-95m")
-model = AutoModelForSeq2SeqLM.from_pretrained("ai-forever/sage-fredt5-distilled-95m")
-model.to("cuda:0")
+model = AutoModelForSeq2SeqLM.from_pretrained("ai-forever/sage-fredt5-distilled-95m", device_map='cuda')
 
 sentence = "И не чсно прохожим в этот день непогожйи почему я веселый такйо"
-text = "<LM>" + sentence
-with torch.inference_mode():
-    encodings = tokenizer(text, max_length=None, padding="longest", truncation=False, return_tensors="pt")
-    for k, v in encodings.items():
-        encodings[k] = v.to("cuda:0")
-    res = model.generate(
-        **encodings,
-        use_cache=True,
-        max_length = encodings["input_ids"].size(1) * 1.5
-    )
-    res = res.cpu().tolist()
-    res = tokenizer.batch_decode(res, skip_special_tokens=True)
-
-print(res)
+inputs = tokenizer(sentence, max_length=None, padding="longest", truncation=False, return_tensors="pt")
+outputs = model.generate(**inputs.to(model.device), max_length = inputs["input_ids"].size(1) * 1.5)
+print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
+
 # ["И не ясно прохожим в этот день непогожий, почему я весёлый такой?"]
+
 ```
 
 ## Limitations