Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
3from IPython.display import display, Math, Latex
4
5model_dir = "kostyabuh21/DistilBART_forLaTeX "
6model = AutoModelForSeq2SeqLM.from_pretrained(model_dir)
7tokenizer = AutoTokenizer.from_pretrained(model_dir)
8
9device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
10model.to(device)
11
12def get_latex(text):
13 inputs = tokenizer(text, return_tensors='pt').to(device)
14 with torch.no_grad():
15 hypotheses = model.generate(
16 **inputs,
17 do_sample=True,
18 top_p=0.95,
19 num_return_sequences=1,
20 repetition_penalty=1.2,
21 max_length=len(text),
22 temperature=0.6,
23 min_length=10,
24 length_penalty=1.0,
25 no_repeat_ngram_size=2
26 )
27 for h in hypotheses:
28 display(Latex(tokenizer.decode(h, skip_special_tokens=True)))
29 print(tokenizer.decode(h, skip_special_tokens=True))
30
31text = 'интеграл от 3 до 5 по икс dx'
32get_latex(text)