Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
3from IPython.display import display, Math, Latex
4
5model_dir = "turnipseason/latext5"
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, num_return_sequences=1,
18 repetition_penalty=1.2,
19 max_length=len(text),
20 num_beams=10,
21 early_stopping=True
22 )
23 for h in hypotheses:
24 display(Latex(tokenizer.decode(h, skip_special_tokens=True)))
25
26text = '''лямбда прописная квадрат минус три равно десять игрек куб
27 При этом шинус икс равен интеграл от экспоненты до трёх игрек штрих'''
28get_latex(text)
29