from tqdm import tqdm
from unsloth import FastLanguageModel
import torch
import json

model_name = "Hktm/llm-jp-3-13b-sft3"

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name=model_name,
    max_seq_length=2048,
    dtype=None,
    load_in_4bit=True,
    token = HF_TOKEN,
)

FastLanguageModel.for_inference(model)

datasets = []
with open("./elyza-tasks-100-TV_0.jsonl", mode="r") as f:
    item = ""
    for line in f:
      line = line.strip()
      item += line
      if item.endswith("}"):
        datasets.append(json.loads(item))
        item = ""


PROMPT_TEMPLATE_WO_DEMO = """### 指示:
下記の質問に回答してください。

### 質問:
{}

### 回答:"""

SPLIT_WORD = "### 回答:"


results = []
for dt in tqdm(datasets):
    input = dt["input"]

    prompt = PROMPT_TEMPLATE_WO_DEMO.format(input).strip()
    inputs = tokenizer([prompt], return_tensors = "pt").to(model.device)
    outputs = model.generate(
        **inputs,
        max_new_tokens = 2048,
        use_cache = True,
        do_sample=False,
        repetition_penalty=1.2
        )
    _pred = tokenizer.decode(outputs[0], skip_special_tokens=True)
    prediction = _pred.split(SPLIT_WORD)[-1].strip()

    result = {
        "task_id": dt["task_id"],
        "input": input,
        "output": prediction,
        # "pred_src": pred_src
        }
    print("\n")
    print(json.dumps(result, ensure_ascii=False, indent=2))

    results.append(result)

with open(
    f"/content/drive/MyDrive/Colab Notebooks/data/20241123_MatsuoLLM_Final/{model_name}_output.jsonl", 
    mode='w',
    encoding='utf-8'
    ) as f:
    for result in results:
        json.dump(result, f, ensure_ascii=False)
        f.write('\n')