import gradio as gr
import os
os.environ["KERAS_BACKEND"] = "tensorflow"                                                                           
import keras
import keras_nlp



gemma_lm = keras_nlp.models.CausalLM.from_preset("hf://sultan-hassan/CosmoGemma_2b_en")

def launch(input):
    template = "Instruction:\n{instruction}\n\nResponse:\n{response}"
    prompt = template.format(
        instruction=input,                                                                   
        response="",
    )
    out = gemma_lm.generate(prompt, max_length=256)
    ind = out.index('Response') + len('Response')+2
    return out[ind:]

iface = gr.Interface(launch,
                     inputs="text",
                     outputs="text")

iface.launch()