llm-tinyllama / app.py
seryeni's picture
Create app.py
7729d35 verified
raw
history blame contribute delete
437 Bytes
from transformers import LlamaTokenizer, pipeline, AutoModelForCausalLM
model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
tokenizer = LlamaTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
prompt = "Who is Ada Lovelace?"
generated_text = generator(prompt, max_length=50)
print(generated_text[0]['generated_text'])