from transformers import AutoTokenizer, TextStreamer
from intel_extension_for_transformers.transformers import AutoModelForCausalLM, WeightOnlyQuantConfig
import time
model_name = "meta-llama/Llama-2-7b-chat-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
streamer = TextStreamer(tokenizer)
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)

while True:
    print("> ", end="")
    prompt = input().strip()
    if prompt == "quit":
        break
    start = time.time()
    b_prompt = "[INST]{}[/INST]".format(prompt)
    inputs = tokenizer(b_prompt, return_tensors="pt").input_ids
    outputs = model.generate(inputs, streamer=streamer,
                num_beams=1, max_new_tokens=512, do_sample=True, repetition_penalty=1.1)
    inference_time = time.time() - start
    print(f'\n> Total Inference Time: {inference_time} seconds')

woq_config = WeightOnlyQuantConfig(compute_dtype="int8", weight_dtype="int4")