app.py

import os
import gradio as gr
from gradio.components import Slider
import torch
from transformers import pipeline

# Model, information and examples ----------------------------------------------
MODEL_NAMES = ["FLOR-1.3B-GL","Cerebras-1.3B-GL"]
model_id = "proxectonos/FLOR-1.3B-GL"
title = "Modelo de xeración de texto FLOR-1.3B-GL"
markdown_description = """
# FLOR-1.3B-GL

💐 **[FLOR-1.3B-GL](https://huggingface.co/proxectonos/FLOR-1.3B-GL)** is a 1.3B parameters multilingual LLM for Galician language.

👀 **Learn more about FLOR-1.3B:** [HF official model card](https://huggingface.co/proxectonos/FLOR-1.3B-GL) and the [Proxecto Nós](https://nos.gal/en/proxecto-nos).
"""

short_prompts_examples = [
    ["A receita tradicional das filloas é"], 
    ["O neno vivía preto de"]
]

few_shot_prompts_examples = [
    ["Responde á seguinte pregunta. \nPregunta: \"Cal é a capital de Noruega? \"\nResposta: \"A capital de Noruega é Oslo.\"\n---- \nResponde á seguinte pregunta.\nPregunta: \"Cal é a moeda de Portugal\" \nResposta: \"A moeda de Portugal é o euro.\" \n---- \nResponde á seguinte pregunta. \nPregunta: \"Cal é a capital de Suecia?\"\nResposta:"],
    ["Extrae as entidades nomeadas do seguinte texto: \nTexto: \"Chámome Wolfgang e vivo en Berlin\" \nEntidades: Wolfgang:PER, Berlin:LOC \n ---- \nExtrae as entidades nomeadas do seguinte texto: \nTexto: \"María e Miguel non teñen ningún problema\" \nEntidades: María:PER, Miguel:PER \n---- \nExtrae as entidades nomeadas do seguinte texto: \nTexto: \"O mellor de Barcelona é o bar do meu amigo Pablo\" \nEntidades: Pablo:PER, Barcelona:LOC \n---- \nExtrae as entidades nomeadas do seguinte texto: \nTexto: \"Carlos comparte cuarto con Marc\" \nEntidades:"],
    ["Cualifica como Positivo ou Negativo o sentimento da seguinte frase:\n Texto: \"Estou moi feliz\"\n Polaridade: Positivo\n ---- \n Cualifica como Positivo ou Negativo o sentimento da seguinte frase:\n Texto: \"Non me gusta beber cervexa\"\n Polaridade: Negativo\n ----\n Cualifica como Positivo ou Negativo o sentimento da seguinte frase:\n Texto: \"O meu pai detesta o seu traballo\"\n Polaridade: Negativo\n ----\n Cualifica como Positivo ou Negativo o sentimento da seguinte frase:\n Texto: \"Uxía desfruta xogando ao fútbol\"\n Polaridade: Positivo\n ----\n Cualifica como Positivo ou Negativo o sentimento da seguinte frase:\n Texto: \"O neno non está contento coas notas\"\n Polaridade:"]
]
fronted_theme = 'Soft'

# Model charge ---------------------------------------------------------
model_id_flor  = "proxectonos/FLOR-1.3B-GL"
generator_model_flor = pipeline("text-generation", model=model_id_flor)
model_id_cerebras  = "proxectonos/Cerebras-1.3B-GL"
generator_model_cerebras = pipeline("text-generation", model=model_id_cerebras, token=os.environ['TOKEN_HF'])

# Generation functions ---------------------------------------------------------
def get_model(model_selection):
    if model_selection == "FLOR-1.3B-GL":
        return generator_model_flor
    else:
        return generator_model_cerebras

def remove_empty_lines(text):
    lines = text.strip().split("\n")
    non_empty_lines = [line for line in lines if line.strip()]
    return "\n".join(non_empty_lines)

def predict(prompt, model_select, max_length, repetition_penalty, temperature):
    print("Dentro da xeración...")
    generator_model = get_model(model_select)
    prompt_length = len(generator_model.tokenizer.encode(prompt))
    generated_text = generator_model(
        prompt,
        max_length=prompt_length + max_length, 
        pad_token_id=generator_model.tokenizer.eos_token_id, 
        repetition_penalty=repetition_penalty,
        temperature=temperature,
        do_sample=True)

    generated_sequence = generated_text[0]['generated_text']
    if generated_sequence is  None:
        gr.Warning('Inference endpoint is not available right now. Please try again later.')
        return
    
    generated_sequence = remove_empty_lines(generated_sequence)
    print("Xeración completada")
    return generated_sequence

# Gradio app ---------------------------------------------------------
def clear(): 
    return (
        None, 
        None,
        gr.update(value=20),
        gr.update(value=1.3),
        gr.update(value=0.5)
    )
def pass_to_input(generated_gl):
    return (
        gr.update(value=generated_gl),
        None
    )

def parameters_default(text):
    return (
        gr.update(value=30), # max_length
        gr.update(value=1.3), # repetition_penalty
        gr.update(value=0.5) # temperature
    )

def parameters_fewshot_prompt(text):
    return (
        gr.update(value=15), # max_length
        gr.update(value=1), # repetition_penalty
        gr.update(value=0.5) # temperature
    )

def gradio_app():
    with gr.Blocks(theme=fronted_theme) as demo:
        with gr.Row():
            with gr.Column(scale=0.1):
                gr.HTML('<img src="https://huggingface.co/spaces/proxectonos/README/resolve/main/title-card.png" width="100%" style="border-radius: 0.75rem;">')
            with gr.Column():
                gr.Markdown(markdown_description)
        with gr.Row(equal_height=True):            
                model_select = gr.Dropdown(
                    label="Escolle un modelo:",
                    choices=MODEL_NAMES,
                    value=MODEL_NAMES[0],
                    interactive=True
                )
        with gr.Row(equal_height=True):
            with gr.Column():
                text_gl = gr.Textbox(label="Input", 
                                     lines=6, placeholder="e.g. O neno vai a escola con ")
                with gr.Row(variant="panel"):
                        with gr.Accordion("Model parameters", open=False):
                            max_length = Slider(
                                minimum=1,
                                maximum=200,
                                step=1,
                                value=30,
                                label="Max tokens"
                            )
                            repetition_penalty = Slider(
                                minimum=0.1,
                                maximum=4,
                                step=0.1,
                                value=1.3,
                                label="Repetition penalty"
                            )
                            temperature = Slider(
                                minimum=0,
                                maximum=1,
                                value=0.5,
                                label="Temperatura"
                            )
                generator_btn = gr.Button(value="Generate",variant='primary')
            with gr.Column():
                generated_gl = gr.Textbox(label="Output", 
                                          lines=6, 
                                          placeholder="Generated text will appear here",
                                          interactive=False,
                                          show_copy_button=True)
                pass_btn = gr.Button(value="Pass text to input")
                clean_btn = gr.Button(value="Clean")

        generator_btn.click(predict, inputs=[text_gl, model_select, max_length, repetition_penalty, temperature], outputs=generated_gl, api_name="generate-flor-gl")
        clean_btn.click(fn=clear, inputs=[], outputs=[text_gl, generated_gl, max_length, repetition_penalty, temperature], queue=False, api_name=False)
        pass_btn.click(fn=pass_to_input, inputs=[generated_gl], outputs=[text_gl,generated_gl], queue=False, api_name=False)
        
        with gr.Row():
            with gr.Column(scale=0.5):
                gr.Examples(
                    label = "Short prompts",
                    examples = short_prompts_examples,
                    inputs = [text_gl],
                    outputs = [max_length, repetition_penalty, temperature],
                    fn = parameters_default,
                    run_on_click = True
                )
                gr.Examples(
                    label = "Few-shot prompts",
                    examples = few_shot_prompts_examples,
                    inputs = [text_gl],
                    outputs = [max_length, repetition_penalty, temperature],
                    fn = parameters_fewshot_prompt,
                    run_on_click = True
                )

    demo.launch()

if __name__ == "__main__":
    gradio_app()