import os
import bitsandbytes as bnb
import pandas as pd
import torch
import torch.nn as nn
import transformers


from peft import (
    LoraConfig,
    PeftConfig,
    get_peft_model,
    prepare_model_for_kbit_training,
    PeftModel
)
from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
)

import gradio as gr

import warnings

warnings.filterwarnings("ignore")
device = "cuda:0"


MODEL_NAME = 'diegi97/dolly-v2-6.9b-sharded-bf16'

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    load_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
)

model =AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    device_map="auto",
    trust_remote_code=True,
    quantization_config=bnb_config,
)

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
tokenizer.pad_token = tokenizer.eos_token

peft_model_id = "AdiOO7/Azure-Classifier-dolly-7B"
# peft_model_id = "SparkExpedition/Ticket-Classifier-dolly-7B"
config = PeftConfig.from_pretrained(peft_model_id)
model = AutoModelForCausalLM.from_pretrained(
    config.base_model_name_or_path,
    return_dict=True,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True,
)
tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
tokenizer.pad_token = tokenizer.eos_token

model = PeftModel.from_pretrained(model, peft_model_id)

generation_config = model.generation_config
generation_config.max_new_tokens = 8
generation_config.num_return_sequences = 1
generation_config.temperature = 0.3
generation_config.top_p = 0.7
generation_config.pad_token_id = tokenizer.eos_token_id
generation_config.eos_token_id = tokenizer.eos_token_id

instruct = "From which azure service the issue is raised from {Power BI/Azure Data Factory/Azure Analysis Services}"

def generate_response(question: str) -> str:

    prompt = f"""
    ### <instruction>: {instruct}
    ### <human>: {question}
    ### <assistant>:
    """.strip()

    encoding = tokenizer(prompt, return_tensors="pt").to(device)
    with torch.inference_mode():
        outputs = model.generate(
            input_ids=encoding.input_ids,
            attention_mask=encoding.attention_mask,
            generation_config=generation_config,
        )
    response = tokenizer.decode(outputs[0],skip_special_tokens=True)

    assistant_start = '<assistant>:'
    response_start = response.find(assistant_start)
    return response[response_start + len(assistant_start):].strip()

labels = ['PowerBI', 'Azure Data Factory', 'Azure Analysis Services']

def answer_prompt(prompt):
  response = generate_response(prompt)
  for lab in labels:
    if response.find(lab) != -1:
      return lab

iface = gr.Interface(fn=answer_prompt,
                     inputs=gr.Textbox(lines=5, label="Enter Your Issue", css={"font-size":"18px"}),
                     outputs=gr.Textbox(lines=5, label="Generated Answer", css={"font-size":"16px"}))

iface.launch()