import gradio as gr
from transformers import pipeline

playground = gr.Blocks()
image_pipe = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")

def launch_image_pipe(input):
    out = image_pipe(input)
    return out[0]['generated_text']

def create_playground_header():
    gr.Markdown("""
                # 🤗 Hugging Face Labs
                **Explore different LLM on Hugging Face platform. Just play and enjoy**
                """)

def create_playground_footer():
    gr.Markdown("""
                **To Learn More about 🤗 Hugging Face, [Click Here](https://huggingface.co/docs)**
                """)

def create_tabs_header(topic, description, references):
    with gr.Row():
        with gr.Column(scale=4):
            # reference_list = "> " + "\n> ".join(references)
            # content  = f"## {topic}\n"
            # content += f"### {description}\n"
            # for ref in references:
            #     content += f"> {ref}\n"
            # gr.Markdown(content)
            gr.Markdown("""
                        ## Image Captioning
                        ### Upload a image, check what AI understand and have vision on it.
                        > category: Image-to-Text
                        > model: [Salesforce/blip-image-captioning-base](https://huggingface.co/Salesforce/blip-image-captioning-base)
                        """)
            
        with gr.Column(scale=1):
            test_pipeline_button = gr.Button(value="Process")
        return test_pipeline_button

with playground:
    create_playground_header()
    with gr.Tabs():
        with gr.TabItem("Image"):
            
            topic = "Image Captioning"
            description = "Upload a image, check what AI understand and have vision on it."
            references = ["category: Image-to-Text",
                            "model: [Salesforce/blip-image-captioning-base](https://huggingface.co/Salesforce/blip-image-captioning-base)"]
            image_pipeline_button = create_tabs_header(topic, description, references)
            
            with gr.Row(visible=True) as use_pipeline:
                with gr.Column():
                    img = gr.Image(type='pil')
                with gr.Column():
                    generated_textbox = gr.Textbox(lines=2, placeholder="", label="Generated Text")
                    
            image_pipeline_button.click(launch_image_pipe,
                                        inputs=[img],
                                        outputs=[generated_textbox])
            
        with gr.TabItem("Text"):
            gr.Markdown("""
                        > Text Summarization and Translation
                        """)
        
        with gr.TabItem("Name Entity"):
            gr.Markdown("""
                        > Name Entity Recognition
                        """)
            
    create_playground_footer()

playground.launch(share=True)