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'

demo = gr.Interface(
    fn=generate_context,
    inputs=[
        gr.Textbox("Once Upon a Time ...", label="Enter Prompt to continue generation"),
        gr.Slider(0.7, 1, value=0.8, label="Temperature", info="Parameter to controls the level of randomness in the sampling process . Higher values increases creativity",step=0.1),
        gr.Slider(100, 300, value=200, label="Top K", info="Limit the response to K next predictions",step=50),
        gr.Slider(10, 100, value=50, label="Max Tokens", info="Set value to get the output to max_values of tokens",step=10),
    ],
    outputs="text",
    examples = examples,
)

demo.launch()