# There is a blog post explaining the results. (Japanese)
# https://webbigdata.jp/post-20104/

from transformers import AutoTokenizer
from auto_gptq import AutoGPTQForCausalLM
from peft import PeftModel, PeftConfig

from transformers import AutoModelForCausalLM
from transformers import GPTQConfig

quantized_model_dir = "./"
model_basename = "model"

tokenizer = AutoTokenizer.from_pretrained(quantized_model_dir, local_files_only=True)

device_map = "auto"
model = AutoModelForCausalLM.from_pretrained(
        quantized_model_dir,
        device_map=device_map,
        use_safetensors=True,
        quantization_config= GPTQConfig(bits=4, disable_exllama=True)
)

adapters_name = f'checkpoint-700'
p_model = PeftModel.from_pretrained(model, adapters_name)

prompt_text = [
"風の谷のナウシカは何歳ですか？",
"風の谷のナウシカに出演するミトは何歳ですか？",
"風の谷のナウシカに出演するユパは何歳ですか？",
"『風の谷のナウシカ』に出演するクシャナはどのような地位にありますか？",
"『風の谷のナウシカ』の王蟲の眼は通常何色で、怒るとどう変化しますか？",
"劇場版アニメ『風の谷のナウシカ』はいつ公開されましたか？",
"漫画『風の谷のナウシカ』全巻の累計発行部数は何部ですか？",
"『風の谷のナウシカ』は何年に歌舞伎化されましたか？",
"『風の谷のナウシカ』で、ナウシカが主に乗る飛行メカは何ですか？",
"『風の谷のナウシカ』に出演するクロトワはどの大学を卒業していますか？"]

for j in range(len(prompt_text)):
	prompt_template = f'以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。\n\n### 指示:\n{prompt_text[j]}\n\n### 応答:'

	tokens = tokenizer(prompt_template, return_tensors="pt").to("cuda:0").input_ids
	output = p_model.generate(input_ids=tokens, max_new_tokens=100, do_sample=True, temperature=0.01)
	print(tokenizer.decode(output[0]))