"""
## Setup and Imports
"""

# import os
# os.environ["KERAS_BACKEND"] = "jax"
import keras_core as keras
import keras_nlp
import gradio as gr
from huggingface_hub import Repository

import keras
from keras.optimizers import Adam
from keras.losses import SparseCategoricalCrossentropy
import keras_nlp

"""## Get the repository"""

repo = Repository(
    local_dir="title-generator-using-summary-gpt2-llm",
    clone_from="Nageswaran/title-generator-using-summary-gpt2-llm",
)

"""## Build the model"""

gpt2_lm = keras.models.load_model(
    "title-generator-using-summary-gpt2-llm/gpt2_lm.keras"
)

title="Summary to Title"
description="GPT2 Model for summary to title."

def get_title(summary):
    summary = summary.replace("\n", " ")
    output = gpt2_lm.generate(f"Summary: {summary} Title:", max_length=600)
    title = output.split("Title:")[1]
    return title

"""## Build the space and launch it"""

gpt_space = gr.Interface(
    fn=get_title,
    inputs=gr.Textbox(label="Summary"),
    outputs=gr.Textbox(label="Title"),
    title=title,
    description=description,
)

gpt_space.launch()