import gradio as gr
from transformers import pipeline

playground = gr.Blocks()
image_pipe = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
get_completion = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")


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

def translate(input_text, source, target):
    try:
      model = f"Helsinki-NLP/opus-mt-{source}-{target}"
      pipe = pipeline("translation", model=model)
      translation = pipe(input_text)
      return translation[0]['translation_text'], ""
    except KeyError:
      return "", f"Error: Translation direction {source_readable} to {target} is not supported by Helsinki Translation Models"

def summarize(input):
    output = get_completion(input)
    summary_origin = output[0]['summary_text']
    summary_translated = translate(summary_origin,'en','fr')
    return summary_origin, summary_translated[0]

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():
                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.Row():
                with gr.Column(scale=4):
                    gr.Markdown("""
                                ## Text Summarization and Translation
                                ### Summarize the paragraph and translate it into other language.
                                > pipeline: summarization, model: [sshleifer/distilbart-cnn-12-6](https://huggingface.co/sshleifer/distilbart-cnn-12-6)
                                > pipeline: translation, model: [Helsinki-NLP/opus-mt-en-fr](https://huggingface.co/Helsinki-NLP/opus-mt-en-fr)
                                """)
                    
                with gr.Column(scale=1):
                    text_pipeline_button = gr.Button(value="Process")

            with gr.Row():
                with gr.Column():
                    source_text = gr.Textbox(label="Text to summarize", lines=6)
                with gr.Column():
                    summary_textbox = gr.Textbox(lines=3, placeholder="", label="Summarization")
                    translated_textbox = gr.Textbox(lines=3, placeholder="", label="Translate Result")

                    
            text_pipeline_button.click(summarize,
                                        inputs=[source_text],
                                        outputs=[summary_textbox, translated_textbox])
            
        with gr.TabItem("Name Entity"):
            gr.Markdown("""
                        > Name Entity Recognition
                        """)
            
    create_playground_footer()

playground.launch(share=True)