| | import torch |
| | from transformers import T5Tokenizer, T5ForConditionalGeneration |
| |
|
| | def pretty_print(text, prompt=True): |
| | s = "" |
| | if prompt: |
| | for section in text.split(', '): |
| | premises = section.split(" and ") |
| | if len(premises) > 1: |
| | for premise in premises[:-1]: |
| | s += premise + "\n\n\n" + "and" + "\n\n\n" |
| | s += premises[-1] + "\n\n\n" |
| | else: |
| | s += section + "\n\n\n" |
| | else: |
| | for equation in text.split("and"): |
| | s += equation + "\n\n\n" |
| | return print(s[:-3]) |
| |
|
| |
|
| | def load_model(model_id): |
| | device = 'cuda' if torch.cuda.is_available() else 'cpu' |
| | tokenizer = T5Tokenizer.from_pretrained(model_id) |
| | model = T5ForConditionalGeneration.from_pretrained(model_id).to(device) |
| | return tokenizer, model |
| |
|
| |
|
| | def inference(prompt, tokenizer, model): |
| | device = 'cuda' if torch.cuda.is_available() else 'cpu' |
| | input_ids = tokenizer.encode(prompt, return_tensors='pt', max_length=512, truncation=True).to(device) |
| | output = model.generate(input_ids=input_ids, max_length=512, early_stopping=True) |
| | generated_text = tokenizer.decode(output[0], skip_special_tokens=True) |
| | |
| | |
| | derivation = generated_text.replace("\\ ","\\") |
| | partial_symbols = derivation.split(" ") |
| | backslash_syms = set([i for i in partial_symbols if "\\" in i]) |
| | for i in range(len(partial_symbols)): |
| | sym = partial_symbols[i] |
| | for b_sym in backslash_syms: |
| | if b_sym.replace("\\","") == sym: |
| | partial_symbols[i] = b_sym |
| | return " ".join(partial_symbols) |
| |
|
