import numpy as np
import pandas as pd
import requests
import os
import gradio as gr
import json
from dotenv import load_dotenv, find_dotenv
_ = load_dotenv(find_dotenv())

from predibase import Predibase, FinetuningConfig, DeploymentConfig

# Get a KEY from https://app.predibase.com/
api_token  = os.getenv('PREDIBASE_API_KEY')
pb = Predibase(api_token=api_token)

adapter_id = 'tour-assistant-model/14'
lorax_client = pb.deployments.client("solar-1-mini-chat-240612")


def extract_json(gen_text, n_shot_learning=0):
    if(n_shot_learning == -1) :
        start_index = 0
    else :
        start_index = gen_text.index("### Response:\n{") + 14
    if(n_shot_learning > 0) :
        for i in range(0, n_shot_learning):
            gen_text = gen_text[start_index:]
            start_index = gen_text.index("### Response:\n{") + 14
    end_index = gen_text.find("}\n\n### ") + 1
    return gen_text[start_index:end_index]

def get_completion(prompt):
    return lorax_client.generate(prompt, adapter_id=adapter_id, max_new_tokens=1000).generated_text

def greet(input):
	total_prompt=f"""
<|im_start|>system\nYou are a helpful support assistant. Answer the following question.<|im_end|>
<|im_start|>question\n How much are union dues, and what do they cover?
<|im_start|>answer\nThe union dues for our union is 3%."<|im_end|>

<|im_start|>system\nYou are a helpful support assistant. Answer the following question.<|im_end|>
<|im_start|>question
{input}. Return as a JSON response<|im_end|>
<|im_start|>answer
"""

	print("***total_prompt:")
	print(total_prompt)
	response = get_completion(total_prompt)
	#gen_text = response["predictions"][0]["generated_text"]
	#return json.dumps(extract_json(gen_text, 3))

	###gen_text = response["choices"][0]["text"]

	#return gen_text

	###return json.dumps(extract_json(gen_text, -1))
	return response

	#return json.dumps(response)

#iface = gr.Interface(fn=greet, inputs="text", outputs="text")
#iface.launch()

#iface = gr.Interface(fn=greet, inputs=[gr.Textbox(label="Text to find entities", lines=2)], outputs=[gr.HighlightedText(label="Text with entities")], title="NER with dslim/bert-base-NER", description="Find entities using the `dslim/bert-base-NER` model under the hood!", allow_flagging="never", examples=["My name is Andrew and I live in California", "My name is Poli and work at HuggingFace"])
iface = gr.Interface(fn=greet, inputs=[gr.Textbox(label="Question", lines=3)], outputs="json")
iface.launch()