Initial commit for Hugging Face Space

.idea/.gitignore

@@ -0,0 +1,3 @@
+# Default ignored files
+/shelf/
+/workspace.xml

.idea/inspectionProfiles/profiles_settings.xml

@@ -0,0 +1,6 @@
+<component name="InspectionProjectProfileManager">
+  <settings>
+    <option name="USE_PROJECT_PROFILE" value="false" />
+    <version value="1.0" />
+  </settings>
+</component>

.idea/modules.xml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectModuleManager">
+    <modules>
+      <module fileurl="file://$PROJECT_DIR$/.idea/multi_beam_text_streamer.iml" filepath="$PROJECT_DIR$/.idea/multi_beam_text_streamer.iml" />
+    </modules>
+  </component>
+</project>

.idea/multi_beam_text_streamer.iml

@@ -0,0 +1,12 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="PYTHON_MODULE" version="4">
+  <component name="NewModuleRootManager">
+    <content url="file://$MODULE_DIR$" />
+    <orderEntry type="inheritedJdk" />
+    <orderEntry type="sourceFolder" forTests="false" />
+  </component>
+  <component name="PyDocumentationSettings">
+    <option name="format" value="PLAIN" />
+    <option name="myDocStringFormat" value="Plain" />
+  </component>
+</module>

.idea/vcs.xml

@@ -0,0 +1,6 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="VcsDirectoryMappings">
+    <mapping directory="" vcs="Git" />
+  </component>
+</project>

Dockerfile

@@ -26,9 +26,5 @@ CMD ["gunicorn", \
      "--worker-class", "eventlet", \
      "--workers", "1", \
      "--timeout", "300", \
-     "--keep-alive", "120", \
-     "--log-level", "debug", \
-     "--worker-connections", "1000", \
-     "--backlog", "2048", \
      "--bind", "0.0.0.0:7860", \
      "app:app"]

app.py

@@ -1,21 +1,13 @@
+import eventlet
+eventlet.monkey_patch()
 from flask import Flask, render_template
 from flask_socketio import SocketIO
 from transformers import MultiBeamTextStreamer, AutoTokenizer, AutoModelForCausalLM
 import torch
 import time
-import eventlet
 
-eventlet.monkey_patch()
 app = Flask(__name__)
-socketio = SocketIO(
-    app,
-    ping_timeout=60,
-    ping_interval=25,
-    cors_allowed_origins="*",
-    async_mode='eventlet',
-    logger=True,
-    engineio_logger=True
-)
+socketio = SocketIO(app, ping_timeout=60)
 
 # Initialize model and tokenizer
 MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
@@ -65,61 +57,57 @@ def index():
 
 @socketio.on('generate')
 def handle_generation(data):
-    def generate_async():
-        try:
-            app.logger.info("Generation started with data: %s", data)
-            socketio.emit('generation_started', callback=lambda: eventlet.sleep(0))
-
-            prompt = data['prompt']
-            num_beams = data.get('num_beams', 5)
-            max_new_tokens = data.get('max_tokens', 512)
-            sleep_time = data.get('sleep_time', 0)
-
-            messages = [
-                {"role": "system", "content": "You are a helpful assistant."},
-                {"role": "user", "content": prompt}
-            ]
-
-            text = tokenizer.apply_chat_template(
-                messages,
-                tokenize=False,
-                add_generation_prompt=True
-            )
-
-            model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
-
-            streamer = WebSocketBeamStreamer(
-                tokenizer=tokenizer,
+    # Emit a generation start event
+    socketio.emit('generation_started')
+
+    prompt = data['prompt']
+    num_beams = data.get('num_beams', 5)
+    max_new_tokens = data.get('max_tokens', 512)
+    sleep_time = data.get('sleep_time', 0)  # Get sleep time from frontend
+
+    # Create messages format
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": prompt}
+    ]
+
+    # Apply chat template
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
+
+    # Prepare inputs
+    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+
+    # Initialize streamer with sleep time
+    streamer = WebSocketBeamStreamer(
+        tokenizer=tokenizer,
+        num_beams=num_beams,
+        sleep_time=sleep_time,
+        skip_prompt=True
+    )
+
+    try:
+        # Generate with beam search
+        with torch.no_grad():
+            model.generate(
+                **model_inputs,
                 num_beams=num_beams,
-                sleep_time=sleep_time,
-                skip_prompt=True
+                num_return_sequences=num_beams,
+                max_new_tokens=max_new_tokens,
+                output_scores=True,
+                return_dict_in_generate=True,
+                early_stopping=True,
+                streamer=streamer
             )
+    except Exception as e:
+        socketio.emit('generation_error', {'error': str(e)})
+    finally:
+        # Emit generation completed event
+        socketio.emit('generation_completed')
 
-            with torch.no_grad():
-                model.generate(
-                    **model_inputs,
-                    num_beams=num_beams,
-                    num_return_sequences=num_beams,
-                    max_new_tokens=max_new_tokens,
-                    output_scores=True,
-                    return_dict_in_generate=True,
-                    early_stopping=True,
-                    streamer=streamer
-                )
-
-        except Exception as e:
-            app.logger.error("Generation error: %s", str(e), exc_info=True)
-            socketio.emit('generation_error', {'error': str(e)})
-        finally:
-            socketio.emit('generation_completed')
-
-    eventlet.spawn(generate_async)
 
 if __name__ == '__main__':
-    socketio.run(
-        app,
-        host='0.0.0.0',
-        port=7860,
-        debug=True,
-        use_reloader=False
-    )
+    socketio.run(app, host='0.0.0.0', port=7860)