import keras
import keras_nlp
import gradio as gr
import os
import re

# Get Secrets
os.environ["KAGGLE_USERNAME"] = os.getenv('KAGGLE_USERNAME')
os.environ["KAGGLE_KEY"] = os.getenv('KAGGLE_KEY')
os.environ["KERAS_BACKEND"] = "jax"

# Establish LLM
gemma_lm = keras_nlp.models.GemmaCausalLM.from_preset("gemma_instruct_2b_en")
sampler = keras_nlp.samplers.TopKSampler(k=1, seed=42)
gemma_lm.compile(sampler=sampler)

# Template
template = """Your task is to enhance user provided prompts for large language models.
Use the following guidelines between backticks to improve prompts.
```
Be specific: Clearly state your desired outcome or the type of information you're seeking. The more specific and focused your prompt is, the better chances you have of obtaining accurate and useful responses.
Provide context: Offer relevant background information to set the stage and ensure that the model understands the context of your request. This to generate more informed and context-aware responses.
Use explicit instructions: If you want the model to perform a specific task, make your instructions explicit. Specify the format, structure, or steps required for the response. For example, instead of asking for a "description," you could write a "step-by-step description" or a "brief summary."
Ask for pros and cons: If you're seeking an evaluation or comparison, explicitly ask for the advantages and disadvantages of a particular topic or approach.
```
Return an enhanced version of the following prompt provided by the Human.
Human:{query} 
"""

# Build Gradio App
def process_input(input_text):
    # Process the input here and generate the output
    prompt = template.format(
    query =input_text
)
    gemma_output = gemma_lm.generate(prompt, max_length=1024)
    output = re.search(r'Enhanced prompt: (.+)', gemma_output).group(1).strip()
    return output

# Create the input interface
input_textbox = gr.Textbox(label="Enter your Prompt")

# Create the output interface
output_textbox = gr.Textbox(label="Enhanced Prompt")

# Create the Gradio app
gr.Interface(
    fn=process_input,
    inputs=input_textbox,
    outputs=output_textbox,
    title="Auto Prompt Engineer",
    description="Enter your prompt and get it evaluated and enchanced for Free!",
    theme="huggingface"
).launch()