from transformers import GPT2Tokenizer, GPT2LMHeadModel, StoppingCriteria, StoppingCriteriaList
import torch
model_path = "iopwsy/MatGPT-synthesis"
tokenizer = GPT2Tokenizer.from_pretrained(model_path, pad_token = '<|endoftext|>')
model = GPT2LMHeadModel.from_pretrained(model_path)
model.config.pad_token_id = model.config.eos_token_id
model.eval()
class StopforGPT2(StoppingCriteria):
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
        return input_ids[0,-1].detach().cpu().numpy() == 50256

def get_path(text):
    with torch.no_grad():
        res = model.generate(input_ids = tokenizer.encode_plus(text, return_tensors = 'pt'),
                             do_sample=True,
                             top_k = 10,
                             top_p = 0.95,
                             temperature = 0.1,
                             max_new_tokens = 300,
                             stopping_criteria=StoppingCriteriaList([StopforGPT2()]))
    print(tokenizer.decode(res[0],skip_special_tokens=True))

### example
get_path("How to synthesis Li7La3Zr2O12?\n")