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

"""
For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
"""
css = """

html, body {
    margin: 0;
    padding: 0;
    height: 100%;
    overflow: hidden;
}
body::before {
    content: '';
    position: fixed;
    top: 0;
    left: 0;
    width: 100vw;
    height: 100vh;
    background-image: url('https://png.pngtree.com/background/20230413/original/pngtree-medical-color-cartoon-blank-background-picture-image_2422159.jpg');
    background-size: cover;
    background-repeat: no-repeat;
    opacity: 0.60;    
    background-position: center;
    z-index: -1;    
}
.gradio-container {
    display: flex;
    flex-direction: column;
    justify-content: center;
    align-items: center;
    height: 100vh;  
}

"""

gemma_model = keras_nlp.models.GemmaCausalLM.from_preset("hf://harishnair04/gemma_instruct_medtr_2b")

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

chat_interface = gr.Interface(
    respond,
    inputs="text",
    outputs="text",
    title="MedTr Chat",
    description="Gemma 2B instruct fine-tuned for medical transcripts",
    css=css
)

chat_interface.launch()