import os
from llama_cpp import Llama
from huggingface_hub import hf_hub_download

from google.colab import userdata
os.environ['HUGGINGFACE_TOKEN'] = userdata.get('HF_TOKEN')

# Download the model
hf_hub_download(
    repo_id="hideax/gemma-2-27b-it-5-q5_k_m",
    filename="unsloth.Q5_K_M.gguf",
    local_dir="./models"
)

llm = Llama(
    model_path="models/unsloth.Q5_K_M.gguf",
    flash_attn=True,
    n_gpu_layers=81,
    n_batch=1024,
    n_ctx=8192,
)

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

from tqdm import tqdm

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

  prompt = f"""### 指示\n{input}\n### 回答\n"""
  output = llm(
            prompt,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            stop=stop or [],
            echo=False
        )
  prediction = output["choices"][0]["text"].strip()
  result = {"task_id": dt["task_id"], "input": input, "output": prediction}
  print(result)
  results.append(result)

model_id = "gemma-2-27b-it-5-Q5_K_M"
with open(f"{model_id}_output.jsonl", 'w', encoding='utf-8') as f:
    for result in results:
        json.dump(result, f, ensure_ascii=False)
        f.write('\n')