Adapt to HF inference endpoints

app.py

@@ -3,31 +3,40 @@ from dotenv import load_dotenv
 import gradio as gr
 from gradio.components import Textbox, Button, Slider, Checkbox
 from AinaTheme import theme
-from sagemaker_endpoint import invoke_endpoint
+from huggingface_hub import InferenceClient
 
 load_dotenv()
 
+def generate(prompt, model_parameters):
+    output = client.text_generation(prompt, **model_parameters, return_full_text=True)
+    return output
+
+
+client = InferenceClient(
+    os.environ.get("HF_INFERENCE_ENDPOINT_URL"),
+    token=os.environ.get("HF_INFERENCE_ENDPOINT_TOKEN")
+)
+
 MAX_NEW_TOKENS = int(os.environ.get("MAX_NEW_TOKENS", default=100))
 MAX_INPUT_CHARACTERS= int(os.environ.get("MAX_INPUT_CHARACTERS", default=100))
 SHOW_MODEL_PARAMETERS_IN_UI = os.environ.get("SHOW_MODEL_PARAMETERS_IN_UI", default=True) == "True"
 
 
-def submit_input(input_, max_new_tokens, repetition_penalty, top_k, top_p, do_sample, num_beams, temperature):
+def submit_input(input_, max_new_tokens, repetition_penalty, top_k, top_p, do_sample, temperature):
     if input_.strip() == "":  
         gr.Warning('Not possible to inference an empty input')
         return None
+    
     model_parameters = {
         "max_new_tokens": max_new_tokens,
         "repetition_penalty": repetition_penalty,
         "top_k": top_k,
         "top_p": top_p,
         "do_sample": do_sample,
-        "num_beams": num_beams,
         "temperature": temperature
     }
      
-    output = invoke_endpoint(input_, model_parameters=model_parameters)
-
+    output = generate(input_, model_parameters)
    
     if output is  None:
         gr.Warning('Inference endpoint is not available right now. Please try again later.')
@@ -42,13 +51,12 @@ def clear():
     return (
         None, 
         None,
-        gr.Slider.update(value=100),
-        gr.Slider.update(value=1.2),
-        gr.Slider.update(value=50),
-        gr.Slider.update(value=0.95),
-        gr.Checkbox.update(value=True),
-        gr.Slider.update(value=4),
-        gr.Slider.update(value=0.5),
+        gr.update(value=MAX_NEW_TOKENS),
+        gr.update(value=1.2),
+        gr.update(value=50),
+        gr.update(value=0.95),
+        gr.update(value=True),
+        gr.update(value=0.5),
     )
     
 def gradio_app():
@@ -117,13 +125,6 @@ def gradio_app():
                             value=True, 
                             label="Do sample"
                         )
-                        num_beams = Slider(
-                            minimum=1, 
-                            maximum=8, 
-                            step=1,
-                            value=4,
-                            label="Beams"
-                        )
                         temperature = Slider(
                             minimum=0, 
                             maximum=1,
@@ -153,7 +154,7 @@ def gradio_app():
                     examples=[
                         ["""La capital de Suècia"""],
                     ],
-                    inputs=[input_, max_new_tokens, repetition_penalty, top_k, top_p, do_sample, num_beams, temperature],
+                    inputs=[input_, max_new_tokens, repetition_penalty, top_k, top_p, do_sample, temperature],
                     outputs=output,
                     fn=submit_input,
                 )
@@ -163,7 +164,7 @@ def gradio_app():
                     examples=[
                         ["Tradueix del Castellà al Català la següent frase: \"Eso es pan comido.\" \nTraducció:"],
                     ],
-                    inputs=[input_, max_new_tokens, repetition_penalty, top_k, top_p, do_sample, num_beams, temperature],
+                    inputs=[input_, max_new_tokens, repetition_penalty, top_k, top_p, do_sample, temperature],
                     outputs=output,
                     fn=submit_input,
                 )
@@ -172,7 +173,7 @@ def gradio_app():
                     examples=[
                         ["""Oració: Els sons melòdics produeixen una sensació de calma i benestar en l'individu. \nParàfrasi: La música és molt relaxant i reconfortant.\n----\nOració: L'animal domèstic mostra una gran alegria i satisfacció. \nParàfrasi: El gos és molt feliç. \n----\nOració: El vehicle es va trencar i vaig haver de contactar amb el servei de remolc perquè el transportés. \nParàfrasi: El cotxe es va trencar i vaig haver de trucar la grua. \n----\nOració: El professor va explicar els conceptes de manera clara i concisa. \nParàfrasi:"""],
                     ],
-                    inputs=[input_, max_new_tokens, repetition_penalty, top_k, top_p, do_sample, num_beams, temperature],
+                    inputs=[input_, max_new_tokens, repetition_penalty, top_k, top_p, do_sample, temperature],
                     outputs=output,
                     fn=submit_input,
                 )
@@ -188,8 +189,8 @@ def gradio_app():
             document.getElementById('inputlenght').style.color =  (i.length > m) ? "#ef4444" : "";
         }""")
     
-        clear_btn.click(fn=clear, inputs=[], outputs=[input_, output, max_new_tokens, repetition_penalty, top_k, top_p, do_sample, num_beams, temperature], queue=False, api_name=False)
-        submit_btn.click(fn=submit_input, inputs=[input_, max_new_tokens, repetition_penalty, top_k, top_p, do_sample, num_beams, temperature], outputs=[output], api_name="get-results")
+        clear_btn.click(fn=clear, inputs=[], outputs=[input_, output, max_new_tokens, repetition_penalty, top_k, top_p, do_sample, temperature], queue=False, api_name=False)
+        submit_btn.click(fn=submit_input, inputs=[input_, max_new_tokens, repetition_penalty, top_k, top_p, do_sample, temperature], outputs=[output], api_name="get-results")
 
         demo.launch(show_api=True)