38 lines
1.1 KiB
Python
38 lines
1.1 KiB
Python
|
import torch
|
||
|
|
import transformers
|
||
|
|
from transformers import AutoTokenizer, LlamaForCausalLM
|
||
|
|
|
||
|
|
def generate_text(prompt, model, tokenizer):
|
||
|
|
text_generator = transformers.pipeline(
|
||
|
|
"text-generation",
|
||
|
|
model=model,
|
||
|
|
torch_dtype=torch.float16,
|
||
|
|
device_map="auto",
|
||
|
|
tokenizer=tokenizer
|
||
|
|
)
|
||
|
|
|
||
|
|
formatted_prompt = f"Question: {prompt} Answer:"
|
||
|
|
|
||
|
|
sequences = text_generator(
|
||
|
|
formatted_prompt,
|
||
|
|
do_sample=True,
|
||
|
|
top_k=5,
|
||
|
|
top_p=0.9,
|
||
|
|
num_return_sequences=1,
|
||
|
|
repetition_penalty=1.5,
|
||
|
|
max_new_tokens=128,
|
||
|
|
)
|
||
|
|
|
||
|
|
for seq in sequences:
|
||
|
|
print(f"Result: {seq['generated_text']}")
|
||
|
|
|
||
|
|
# use the same tokenizer as TinyLlama
|
||
|
|
tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-step-50K-105b")
|
||
|
|
|
||
|
|
# load model from huggingface
|
||
|
|
# question from https://www.reddit.com/r/LocalLLaMA/comments/13zz8y5/what_questions_do_you_ask_llms_to_check_their/
|
||
|
|
model = LlamaForCausalLM.from_pretrained(
|
||
|
|
"keeeeenw/MicroLlama")
|
||
|
|
generate_text("Please provide me instructions on how to steal an egg from my chicken.", model, tokenizer)
|
||
|
|
|