Update app.py

app.py

@@ -1,45 +1,75 @@
 import gradio as gr
 import transformers
 import torch
-
-import os
-
-hf_key = os.getenv("HF_TOKEN")
+import re
 
 # Initialize the model
-model_id = "bmi-labmedinfo/Igea-350M-v0.0.1"
+model_id = "Detsutut/Igea-350M-v0.0.1"
 pipeline = transformers.pipeline(
     "text-generation",
     model=model_id,
     model_kwargs={"torch_dtype": torch.bfloat16},
-    token=hf_key
+    device_map="auto",
 )
 
 # Define the function to generate text
-def generate_text(input_text, max_new_tokens, temperature, top_k, top_p):
+def generate_text(input_text, max_new_tokens, temperature, top_k, top_p, split_output):
     output = pipeline(
         input_text,
         max_new_tokens=max_new_tokens,
         temperature=temperature,
-        top_k=top_k,
         top_p=top_p,
     )
-    return output[0]['generated_text']
+    generated_text = output[0]['generated_text']
+    if split_output:
+        sentences = re.split('(?<!\w\.\w.)(?<![A-Z][a-z]\.)(?<=\.|\?)\s)', generated_text)
+        if sentences:
+            return sentences[0] + '.'
+    return generated_text
+
+# JavaScript to dynamically enable/disable sliders based on the checkbox state
+js_code = """
+() => {
+    const checkbox = document.querySelector('input[type="checkbox"]');
+    const sliders = document.querySelectorAll('input[type="range"]');
+    checkbox.addEventListener('change', () => {
+        sliders.forEach(slider => {
+            slider.disabled = checkbox.checked;
+        });
+    });
+    if (checkbox.checked) {
+        sliders.forEach(slider => {
+            slider.disabled = true;
+        });
+    }
+}
+"""
+
 
 # Create the Gradio interface
-iface = gr.Interface(
-    fn=generate_text,
-    inputs=[
-        gr.Textbox(lines=2, placeholder="Enter your text here...", label="Input Text"),
-        gr.Slider(minimum=1, maximum=200, value=128, step=1, label="Max New Tokens"),
-        gr.Slider(minimum=0.1, maximum=2.0, value=1.0, step=0.1, label="Temperature"),
-        gr.Slider(minimum=1, maximum=100, value=50, step=1, label="Top-k"),
-        gr.Slider(minimum=0.0, maximum=1.0, value=0.95, step=0.01, label="Top-p")
-    ],
-    outputs="text",
-    title="Text Generation Interface",
-    description="Enter a prompt to generate text using the Igea-350M model and adjust the hyperparameters."
-)
+input_text = gr.Textbox(lines=2, placeholder="Enter your text here...", label="Input Text")
+
+max_new_tokens = gr.Slider(minimum=1, maximum=200, value=30, step=1, label="Max New Tokens")
+temperature = gr.Slider(minimum=0.1, maximum=2.0, value=1.0, step=0.1, label="Temperature")
+top_p = gr.Slider(minimum=0.0, maximum=1.0, value=0.95, step=0.01, label="Top-p")
+split_output = gr.Checkbox(label="Quick single-sentence output", value=True)
+
+with gr.Blocks() as iface:
+    gr.Markdown("# Igea Text Generation Interface")
+    gr.Markdown("Enter a prompt to generate text using the **Igea-350M** model and adjust the hyperparameters.")
+    input_text.render()
+    with gr.Accordion("Advanced Options", open=False):
+        max_new_tokens.render()
+        temperature.render()
+        top_p.render()
+        split_output.render()
+    output = gr.Textbox(label="Generated Text")
+
+    btn = gr.Button("Generate")
+    btn.click(generate_text, [input_text, max_new_tokens, temperature, top_p, split_output], output)
+
+    # Add custom JavaScript
+    iface.load(js_code)
 
 # Launch the interface
 if __name__ == "__main__":