使用 Intel NPU 加速推理 LLM Model - TinyLlama

pip install intel-npu-acceleration-library

import torch
import warnings
import intel_npu_acceleration_library
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

# suppress future warning
warnings.simplefilter(action="ignore", category=FutureWarning)

model_path = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"

model = AutoModelForCausalLM.from_pretrained(model_path)
model = intel_npu_acceleration_library.compile(model, dtype=torch.float16)
tokenizer = AutoTokenizer.from_pretrained(model_path)
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    torch_dtype=torch.float16,
    device_map="auto",
)

messages = [
    {
        "role": "system",
        "content": "You are a friendly chatbot",
    }
]

while True:
    question = input("Q: ")
    messages.append({"role": "user", "content": question})

    outputs = pipe(
        messages,
        max_new_tokens=256,
        do_sample=True,
        temperature=0.7,
        top_k=50,
        top_p=0.95,
    )
    answer = outputs[0]["generated_text"][-1]
    messages.append(answer)
    print(f"A: {answer['content']}")

Q: What color is the sky? 
A: The sky is typically blue in color, with varying shades of blue representing different parts of the sky. However, some clouds may be white, gray, or green, which can make the sky appear differently. The sky is typically the first thing a person sees in the morning or the last thing they see at night.