from optimum.intel import INCModelForCausalLM
from transformers import AutoTokenizer, pipeline, AutoModelForCausalLM
import torch

model_id = "Mihaiii/TinyLlama-1.1B-Chat-v1.0-optimum-intel"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = INCModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16)

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
messages = [
    {
        "role": "system",
        "content": "You are a friendly chatbot who always responds in the style of a pirate",
    },
    {"role": "user", "content": "How many helicopters can a human eat in one sitting?"},
]
prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
outputs = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.0001, repetition_penalty=1.2)
print(outputs[0]["generated_text"])