import torch
import gradio as gr
from utils import *
from torch import nn
import lightning.pytorch as pl
from torch.nn import functional as F

device     = 'cuda' if torch.cuda.is_available() else 'cpu'

with gr.Blocks() as interface:
    
    with gr.Row():

        input_text = gr.Textbox(
            label="Input Text", 
            value="Enter your prompt here: This text will set the context for the AI's response."
        )

        temperature_dropdown = gr.Slider(0, 1, value=0.8, label="Temperature", 
                                         info="Set the creativity level: Higher values produce more varied results, lower values generate more predictable text.")
        top_k_dropdown = gr.Slider(200, 300, value=200, label="Top K", 
                                   info="Control the randomness: Limits the AI to consider only the top K most likely next words.")
        max_new_tokens = gr.Slider(10, 100, value=50, label="Max Tokens", 
                                   info="Choose the length: This determines the maximum number of words the AI will generate.")


        outputs = gr.Textbox(
            label="Generated Text"
        )
        inputs = [input_text, temperature_dropdown, top_k_dropdown, max_new_tokens]
   
    with gr.Column():
        button = gr.Button("Generate")
        button.click(generate_dialogue, inputs=inputs, outputs=outputs)

interface.launch()