import os
import re
from datetime import datetime
import gradio as gr
import json
from dotenv import load_dotenv, find_dotenv
_ = load_dotenv(find_dotenv())

from training.consts import DEFAULT_INPUT_MODEL, SUGGESTED_INPUT_MODELS
from training.trainer import load_training_dataset, load_tokenizer
from training.generate import generate_response, load_model_tokenizer_for_generate

gpu_family = "a100"

model_dir = "model"

model, tokenizer = load_model_tokenizer_for_generate(model_dir)

def get_completion(prompt, model="dolly-v0-70m"):
	# Examples from https://www.databricks.com/blog/2023/03/24/hello-dolly-democratizing-magic-chatgpt-open-models.html
	instructions = [prompt]
	# set some additional pipeline args
	pipeline_kwargs = {'torch_dtype': "auto"}
	#if gpu_family == "v100":
	#pipeline_kwargs['torch_dtype'] = "float16"
	#elif gpu_family == "a10" or gpu_family == "a100":
	#pipeline_kwargs['torch_dtype'] = "bfloat16"

	pipeline_kwargs['max_new_tokens'] = 300

	# Use the model to generate responses for each of the instructions above.
	for instruction in instructions:
		response = generate_response(instruction, model=model, tokenizer=tokenizer, **pipeline_kwargs)
	if response:
		print(f"Instruction: {instruction}\n\n{response}\n\n-----------\n")
	return response

def greet(input):
	prompt = f"""
Text: ```{input}```
    """
	response = get_completion(prompt)
	return 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="Prompt")], outputs="text")
iface.launch()