from peft import PeftModel, PeftConfig
from transformers import AutoModelForCausalLM,AutoTokenizer

model_id = "meta-llama/Llama-2-13b-chat-hf"
tokenizer = AutoTokenizer.from_pretrained(model_id)
tokenizer.pad_token = tokenizer.eos_token

config = PeftConfig.from_pretrained("zera09/llama_FT")
base_model = AutoModelForCausalLM.from_pretrained(model_id,load_in_4bit=True, device_map='cuda')
model = PeftModel.from_pretrained(base_model, "zera09/llama_FT")

template = """### Instruction
Given this context: {context} and price:{price}output onle one decision from the square bracket [buy,sell,hold] and provide reasoning om why.

### Response:
Decision:
Reasonong:
```"""

from transformers import set_seed


def gen(text):
    toks = tokenizer(text, return_tensors="pt").to("cuda")

    set_seed(32)
    model.eval()
    with torch.no_grad():
        out = model.generate(
            **toks,
            max_new_tokens=350,
            top_k=5,
            do_sample=True,
        )
    return tokenizer.decode(
        out[0][len(toks["input_ids"][0]) :], skip_special_tokens=True
    )

context = "The global recloser control market is expected to grow significantly, driven by increasing demand for power quality and reliability, especially in the electric segment and emerging economies like China. The positive score for this news is 1.1491235518690246e-08. The neutral score for this news is 0.9999998807907104. The negative score for this news is 6.358970239261907e-08"
price = str(12.1)
print(gen(template.format(context=news,price=price)).split("```"))

import pandas as pd

data = panda.read_pickle('./DRIV_train.pkl')
data = pd.DataFrame(data).T

model.eval()
answer_list = []
for idx,row in ans_sum.iterrows():
  toks = tokenizer(template.format(context=row['news']['DRIV'][0],price=str(row['price']['DRIV'][0]))), return_tensors="pt").to("cuda")
    with torch.no_grad():
      out = model.generate(
        **toks,
        max_new_tokens=350,
        top_k=5,
        do_sample=True,
      )
        
    ans_list.append(tokenizer.decode(
                out[0][len(toks["input_ids"][0]) :], skip_special_tokens=True
                )