Spaces:
Sleeping
Sleeping
Added EOS toke stuff increased new token limit and added QOL features to frontent
8e724ea
from transformers import AutoTokenizer, AutoModelForCausalLM | |
model_name = "roneneldan/TinyStories-1M" | |
def load(): | |
global model | |
global tokenizer | |
model = AutoModelForCausalLM.from_pretrained(model_name) | |
tokenizer = AutoTokenizer.from_pretrained(model_name) | |
def generate(input_text): | |
input_ids = tokenizer.encode(input_text, return_tensors="pt") | |
output_ids = model.generate( | |
input_ids, | |
no_repeat_ngram_size=2, | |
max_new_tokens=200, | |
eos_token_id=tokenizer.eos_token_id, | |
temperature=0.2 | |
) | |
return tokenizer.decode(output_ids[0], skip_special_tokens=True) |