import torch, sys
from peft import PeftModel
from transformers import AutoTokenizer, AutoModelForCausalLM

import warnings
warnings.filterwarnings("ignore")

access_token = sys.argv[2]
device = "xpu:0" if sys.argv[1] == "gpu" else "cpu:0"

tokenizer = AutoTokenizer.from_pretrained("./tokenizer/")
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

base_model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-2b",
    token=access_token,
    low_cpu_mem_usage=True,
    return_dict=True,
    torch_dtype=torch.bfloat16,
)

model = PeftModel.from_pretrained(base_model, "adapter_model")
model = model.to(device)

print("Prompt:", " ".join(sys.argv[3:]))

inputs = tokenizer(" ".join(sys.argv[3:]), return_tensors="pt").to(device)
outputs = model.generate(**inputs, max_new_tokens=200, 
                            do_sample=False, top_k=100,temperature=0.1, 
                            eos_token_id=tokenizer.eos_token_id)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))