ai-forever commited on
Commit
5a7a996
1 Parent(s): 2bebc33

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +7 -18
README.md CHANGED
@@ -165,26 +165,15 @@ We compare our solution with both open automatic spell checkers and the ChatGPT
165
  from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
166
 
167
  tokenizer = AutoTokenizer.from_pretrained("ai-forever/sage-fredt5-large")
168
- model = AutoModelForSeq2SeqLM.from_pretrained("ai-forever/sage-fredt5-large")
169
-
170
- model.to("cuda:0")
171
 
172
  sentence = "И не чсно прохожим в этот день непогожйи почему я веселый такйо"
173
- text = "<LM>" + sentence
174
- with torch.inference_mode():
175
- encodings = tokenizer(text, max_length=None, padding="longest", truncation=False, return_tensors="pt")
176
- for k, v in encodings.items():
177
- encodings[k] = v.to("cuda:0")
178
- res = model.generate(
179
- **encodings,
180
- use_cache=True,
181
- max_length = encodings["input_ids"].size(1) * 1.5
182
- )
183
- res = res.cpu().tolist()
184
- res = tokenizer.batch_decode(res, skip_special_tokens=True)
185
- print(res)
186
-
187
- # ["И не ясно прохожим в этот день непогожий, почему я веселый такой."]
188
  ```
189
 
190
  ## Limitations
 
165
  from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
166
 
167
  tokenizer = AutoTokenizer.from_pretrained("ai-forever/sage-fredt5-large")
168
+ model = AutoModelForSeq2SeqLM.from_pretrained("ai-forever/sage-fredt5-large", device_map='cuda')
 
 
169
 
170
  sentence = "И не чсно прохожим в этот день непогожйи почему я веселый такйо"
171
+ inputs = tokenizer(sentence, max_length=None, padding="longest", truncation=False, return_tensors="pt")
172
+ outputs = model.generate(**inputs.to(model.device), max_length = inputs["input_ids"].size(1) * 1.5)
173
+ print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
174
+
175
+ # [ не ясно прохожим в этот день непогожий, почему я весёлый такой?"]
176
+
 
 
 
 
 
 
 
 
 
177
  ```
178
 
179
  ## Limitations