Interactive Llama 2

interactive_llama2.py

from transformers import AutoTokenizer
import transformers
import torch

model = "meta-llama/Llama-2-7b-chat-hf"

print("AutoTokenizer.from_pretrained(model)")

tokenizer = AutoTokenizer.from_pretrained(model)

print("pipeline = transformers.pipeline(...")

pipeline = transformers.pipeline(
    "text-generation",
    model=model,
    torch_dtype=torch.float16,
    device_map="auto",
)

print("sequences = pipeline(...")
print("Entering interactive mode...")

while True:
    sequences = pipeline(
        input("prompt: ") + "\n",
        do_sample=True,
        top_k=10,
        num_return_sequences=1,
        eos_token_id=tokenizer.eos_token_id,
        # max_length=200,
    )

    for seq in sequences:
        print(seq['generated_text'])