from __future__ import annotations
from typing import Iterable
import gradio as gr
from gradio.themes.base import Base
from gradio.themes.utils import colors, fonts, sizes

from llama_cpp import Llama
from huggingface_hub import hf_hub_download

hf_hub_download(repo_id="TheBloke/dolphin-2_6-phi-2-GGUF", filename="dolphin-2_6-phi-2.Q5_K_S.gguf", local_dir=".")
llm = Llama(model_path="./dolphin-2_6-phi-2.Q5_K_S.gguf")

ins = '''<|im_start|>user
{question}<|im_end|>
<|im_start|>assistant
'''

theme = gr.themes.Monochrome(
	primary_hue="purple",
	secondary_hue="blue",
	neutral_hue="neutral",
	radius_size=gr.themes.sizes.radius_sm,
	font=[gr.themes.GoogleFont("Space Grotesk"), "ui-sans-serif", "system-ui", "sans-serif"],
)

def generate(instruction):
	prompt = ins.format(question=instruction)
	response = llm(prompt, stop=['<|im_start|>user', '<|im_end|>'])
	result = response['choices'][0]['text']
	return result

examples = [
	"How do dogs bark?",
	"Why are apples red?",
	"How do I make a campfire?",
	"Why do cats love to chirp at something?"
]

def process_example(args):
	for x in generate(args):
		pass
	return x

css = ".generating {visibility: hidden}"

class BlueTheme(Base):
	def __init__(
		self,
		*,
		primary_hue: colors.Color | str = colors.purple,
		secondary_hue: colors.Color | str = colors.blue,
		neutral_hue: colors.Color | str = colors.neutral,
		spacing_size: sizes.Size | str = sizes.spacing_md,
		radius_size: sizes.Size | str = sizes.radius_md,
		font: fonts.Font
		| str
		| Iterable[fonts.Font | str] = (
			fonts.GoogleFont("Space Grotesk"),
			"ui-sans-serif",
			"sans-serif",
		),
		font_mono: fonts.Font
		| str
		| Iterable[fonts.Font | str] = (
			fonts.GoogleFont("Space Mono"),
			"ui-monospace",
			"monospace",
		),
	):
		super().__init__(
			primary_hue=primary_hue,
			secondary_hue=secondary_hue,
			neutral_hue=neutral_hue,
			spacing_size=spacing_size,
			radius_size=radius_size,
			font=font,
			font_mono=font_mono,
		)
		super().set(
			button_primary_background_fill="linear-gradient(90deg, *primary_300, *secondary_400)",
			button_primary_background_fill_hover="linear-gradient(90deg, *primary_200, *secondary_300)",
			button_primary_text_color="white",
			button_primary_background_fill_dark="linear-gradient(90deg, *primary_600, *secondary_800)",
			block_shadow="*shadow_drop_lg",
			button_shadow="*shadow_drop_lg",
			input_background_fill="zinc",
			input_border_color="*secondary_300",
			input_shadow="*shadow_drop",
			input_shadow_focus="*shadow_drop_lg",
		)


custom_theme = BlueTheme()

with gr.Blocks(theme=custom_theme, analytics_enabled=False, css=css) as demo:
	with gr.Column():
		gr.Markdown(
			""" # 𓄿 [Corvus] Dolphin-2.6 (Phi-2)
			Type in the box below and click the button to generate answers to your most pressing questions!
		""")

		with gr.Row():
			with gr.Column(scale=3):
				instruction = gr.Textbox(placeholder="Enter your question here", label="Question", elem_id="q-input")

				with gr.Box():
					gr.Markdown("**Answer**")
					output = gr.Markdown(elem_id="q-output")
				submit = gr.Button("Generate", variant="primary")
				gr.Examples(
					examples=examples,
					inputs=[instruction],
					cache_examples=False,
					fn=process_example,
					outputs=[output],
				)

	submit.click(generate, inputs=[instruction], outputs=[output])
	instruction.submit(generate, inputs=[instruction], outputs=[output])

demo.queue(concurrency_count=1).launch(debug=True)