Update app.py

app.py

@@ -1,13 +1,7 @@
-# Imports
 import gradio as gr
 import os
 import random
-import time
-import spaces
-import torch
 import threading
-import socket
-import requests
 
 from llama_cpp import Llama
 from huggingface_hub import hf_hub_download
@@ -15,124 +9,92 @@ from huggingface_hub import hf_hub_download
 # Variables
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
-REPO = "bartowski/Mistral-Small-Instruct-2409-GGUF"
-FILE = "Mistral-Small-Instruct-2409-Q6_K_L.gguf"
+REPO = "HuggingFaceTB/SmolLM2-1.7B-Instruct-GGUF"
+FILE = "smollm2-1.7b-instruct-q4_k_m.gguf"
 
 TIMEOUT = 30
 
 MAX_SEED = 9007199254740991
-DEFAULT_INPUT = ""
-DEFAULT_SYSTEM = "Statical is a totally normal human named 'Statical'. Statical responds very briefly, truthfully, and accurately."
 
-TAG_USER = "USER"
-TAG_ASSISTANT = "STATICAL"
+model = Llama(
+    model_path=hf_hub_download(repo_id=REPO, filename=FILE, token=HF_TOKEN),
+    n_ctx=32768,
+    n_threads=48,
+    n_batch=512,
+    n_gpu_layers=0,
+    verbose=True
+)
 
-DEFAULT_SEPARATOR = ","
-DEFAULT_STOP_SEQUENCES = f"{TAG_USER}:,{TAG_ASSISTANT}:,</s>"
-
-model = Llama(model_path=hf_hub_download(repo_id=REPO, filename=FILE, token=HF_TOKEN), n_ctx=32768, n_threads=48, n_batch=512, n_gpu_layers=0, verbose=True) 
-
-# Functions
 def get_seed(seed):
-    seed = seed.strip()
-    if seed.isdigit():
-        return int(seed)
+    if seed and seed.strip().isdigit():
+        return int(seed.strip())
     else:
         return random.randint(0, MAX_SEED)
 
-def generate(input=DEFAULT_INPUT, history=[], system=DEFAULT_SYSTEM, stream=False, temperature=1, top_p=0.95, top_k=50, rep_p=1.2, max_tokens=64, seed=None, separator=DEFAULT_SEPARATOR, stop_sequences=DEFAULT_STOP_SEQUENCES):
+def generate(prompt, temperature, top_p, top_k, repetition_penalty, max_tokens, seed):
     print("[GENERATE] Model is generating...")
-    
-    memory = ""
-    for item in history:
-        if item[0]:
-            memory += f"{TAG_USER}: {item[0].strip()}\n"
-        if item[1]:
-            memory += f"{TAG_ASSISTANT}: {item[1].strip()}</s>\n"
-    prompt = f"{system.strip()}\n{memory}{TAG_USER}: {input.strip()}\n{TAG_ASSISTANT}: "
-    
-    print(prompt)
-    
+
     parameters = {
         "prompt": prompt,
         "temperature": temperature,
         "top_p": top_p,
-        "top_k": top_k,
-        "repeat_penalty": rep_p,
-        "max_tokens": max_tokens,
-        "stop": [seq.strip() for seq in stop_sequences.split(separator)] if stop_sequences else [],
+        "top_k": int(top_k),
+        "repeat_penalty": repetition_penalty,
+        "max_tokens": int(max_tokens),
         "seed": get_seed(seed),
-        "stream": stream
+        "stream": True
     }
-    
+
+    print("Parameters:", parameters)
+
     event = threading.Event()
+    timer = threading.Timer(TIMEOUT, event.set)
+    timer.start()
 
     try:
         output = model.create_completion(**parameters)
         print("[GENERATE] Model has generated.")
-        if stream:
-            buffer = ""
-            timer = threading.Timer(TIMEOUT, event.set)
-            timer.start()
-            try:
-                for _, item in enumerate(output):
-                    if event.is_set():
-                        raise TimeoutError("[ERROR] Generation timed out.")
-                    buffer += item["choices"][0]["text"]
-                    yield buffer
-                    timer.cancel()
-                    timer = threading.Timer(TIMEOUT, event.set)
-                    timer.start()
-            finally:
-                timer.cancel()
-        else:
-            yield output["choices"][0]["text"]
+        buffer = ""
+        try:
+            for _, item in enumerate(output):
+                if event.is_set():
+                    raise TimeoutError("[ERROR] Generation timed out.")
+                buffer += item["choices"][0]["text"]
+                yield buffer
+        finally:
+            timer.cancel()
     except TimeoutError as e:
         yield str(e)
     finally:
         timer.cancel()
 
-@spaces.GPU(duration=15)
-def gpu():
-    return
-    
 # Initialize
-theme = gr.themes.Default(
-    primary_hue="violet",
-    secondary_hue="indigo",
-    neutral_hue="zinc",
-    spacing_size="sm",
-    radius_size="lg",
-    font=[gr.themes.GoogleFont('Kanit'), 'ui-sans-serif', 'system-ui', 'sans-serif'],
-    font_mono=[gr.themes.GoogleFont('Kanit'), 'ui-monospace', 'Consolas', 'monospace'],
-).set(background_fill_primary='*neutral_50', background_fill_secondary='*neutral_100')
-
-model_base = "https://huggingface.co/MaziyarPanahi/WizardLM-2-8x22B-GGUF" # [::-1]
-model_quant = "https://huggingface.co/alpindale/WizardLM-2-8x22B" # [::-1]
-
-with gr.Blocks(theme=theme) as main:
-    with gr.Column():
-        gr.Markdown("# 👁️‍🗨️ WizardLM")
-        gr.Markdown("⠀⠀• ⚡ A text generation inference for one of the best open-source text models: WizardLM-2-8x22B.")
-        gr.Markdown("⠀⠀• ⚠️ WARNING! The inference is very slow due to the model being HUGE; it takes 10 seconds before it starts generating; please avoid high max token parameters and sending large amounts of text; note it uses CPU because I cannot figure out how to run it in GPU without overloading the model.")
-        gr.Markdown(f"⠀⠀• 🔗 Link to models: {model_base} (BASE), {model_quant} (QUANT)")
-
-    with gr.Column():
-        gr.ChatInterface(
-            fn=generate,
-            additional_inputs_accordion=gr.Accordion(label="⚙️ Configurations", open=False, render=False),
-            additional_inputs=[
-                gr.Textbox(lines=1, value=DEFAULT_SYSTEM, label="🪄 System", render=False),
-                gr.Checkbox(label="⚡ Stream", value=True, render=False),
-                gr.Slider(minimum=0, maximum=2, step=0.01, value=1, label="🌡️ Temperature", render=False),
-                gr.Slider(minimum=0.01, maximum=0.99, step=0.01, value=0.95, label="🧲 Top P", render=False),
-                gr.Slider(minimum=1, maximum=2048, step=1, value=50, label="📊 Top K", render=False),
-                gr.Slider(minimum=0.01, maximum=2, step=0.01, value=1.2, label="📚 Repetition Penalty", render=False),
-                gr.Slider(minimum=1, maximum=2048, step=1, value=256, label="⏳ Max New Tokens", render=False),
-                gr.Textbox(lines=1, value="", label="🌱 Seed (Blank for random)", render=False),
-                gr.Textbox(lines=1, value=DEFAULT_SEPARATOR, label="🏷️ Stop Sequences Separator", render=False),
-                gr.Textbox(lines=1, value=DEFAULT_STOP_SEQUENCES, label="🛑 Stop Sequences (Blank for none)", render=False),
-            ]
-        )
-
-main.launch(show_api=False)
+model_base = "Any"
+model_quant = "Any Quant"
+
+with gr.Blocks() as demo:
+    gr.Markdown("# 👁️‍🗨️ WizardLM")
+    gr.Markdown("• ⚡ A text generation inference for one of the best open-source text models: WizardLM-2-8x22B.")
+    gr.Markdown("• ⚠️ **WARNING!** The inference is very slow due to the model being **HUGE**; it takes about 10 seconds before it starts generating. Please avoid high max token parameters and sending large amounts of text. Note it uses CPU because running it on GPU overloads the model.")
+    gr.Markdown(f"• 🔗 Link to models: [{model_base}]({model_base}) (BASE), [{model_quant}]({model_quant}) (QUANT)")
+
+    prompt = gr.Textbox(lines=4, label="Enter your prompt")
+    output = gr.Textbox(lines=10, label="Model output")
+
+    with gr.Accordion("⚙️ Configurations", open=False):
+        temperature = gr.Slider(minimum=0.0, maximum=2.0, value=1.0, step=0.01, label="🌡️ Temperature")
+        top_p = gr.Slider(minimum=0.0, maximum=1.0, value=0.95, step=0.01, label="🧲 Top P")
+        top_k = gr.Slider(minimum=1, maximum=2048, value=50, step=1, label="📊 Top K")
+        repetition_penalty = gr.Slider(minimum=0.0, maximum=2.0, value=1.2, step=0.01, label="📚 Repetition Penalty")
+        max_tokens = gr.Slider(minimum=1, maximum=2048, value=256, step=1, label="⏳ Max New Tokens")
+        seed = gr.Textbox(lines=1, label="🌱 Seed (Blank for random)", value="")
+
+    generate_button = gr.Button("Generate")
+
+    generate_button.click(
+        fn=generate,
+        inputs=[prompt, temperature, top_p, top_k, repetition_penalty, max_tokens, seed],
+        outputs=output,
+    )
+
+demo.launch()