Add Cerebras model

app.py

@@ -1,9 +1,11 @@
+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 = """
@@ -27,17 +29,26 @@ few_shot_prompts_examples = [
 fronted_theme = 'Soft'
 
 # Model charge ---------------------------------------------------------
-model_id  = "proxectonos/FLOR-1.3B-GL"
-generator_model = pipeline("text-generation", model=model_id)
+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, max_length, repetition_penalty, temperature):
+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,
@@ -91,7 +102,13 @@ def gradio_app():
                 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():            
+                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", 
@@ -128,7 +145,7 @@ def gradio_app():
                 pass_btn = gr.Button(value="Pass text to input")
                 clean_btn = gr.Button(value="Clean")
 
-        generator_btn.click(predict, inputs=[text_gl,max_length, repetition_penalty, temperature], outputs=generated_gl, api_name="generate-flor-gl")
+        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)