BeveledCube/bevelapi

1from transformers import AutoTokenizer, AutoModelForCausalLM2 3model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"4 5def load():6  global model7  global tokenizer8  9  model = AutoModelForCausalLM.from_pretrained(model_name)10  tokenizer = AutoTokenizer.from_pretrained(model_name)11 12def generate(input_text):13  input_ids = tokenizer.encode(input_text, return_tensors="pt")14  output_ids = model.generate(input_ids, no_repeat_ngram_size=2, max_new_tokens=200, num_beams=2, eos_token_id=tokenizer.eos_token_id)15  16  return tokenizer.decode(output_ids[0], skip_special_tokens=True)