PhelixZhen/

Algae-550M

from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
import torch

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

tokenizer = AutoTokenizer.from_pretrained('pathtotokenizer')
model = AutoModelForCausalLM.from_pretrained('pathtomodel').to(device)
tokenizer.pad_token = tokenizer.eos_token

txt = 'inputtext'

# greedy search
gen_conf = GenerationConfig(
    num_beams=1,
    do_sample=True,
    max_length=700,
    no_repeat_ngram_size=6,
    eos_token_id=tokenizer.eos_token_id,
    pad_token_id=tokenizer.pad_token_id,
    temperature=0.93,
    top_k=36,
    top_p=0.80
)

tokend = tokenizer.encode_plus(text=txt)
input_ids, attention_mask = torch.LongTensor([tokend.input_ids]).to(device), \
    torch.LongTensor([tokend.attention_mask]).to(device)

outputs = model.generate(
    inputs=input_ids,
    attention_mask=attention_mask,
    generation_config=gen_conf,
    
)

outs = tokenizer.decode(outputs[0].cpu().numpy(), clean_up_tokenization_spaces=True,)
print(outs)