Initial commit for Hugging Face Space

Dockerfile

@@ -9,8 +9,11 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
 
 RUN mkdir -p /app/cache && chmod -R 777 /app/cache
 ENV HF_HOME=/app/cache
+
 ENV PYTHONUNBUFFERED=1
 ENV EVENTLET_NO_GREENDNS=yes
+ENV EVENTLET_THREADPOOL_SIZE=32
+ENV EVENTLET_WEBSOCKET_MONITOR_TIMEOUT=60
 
 COPY . /app
 
@@ -25,5 +28,7 @@ CMD ["gunicorn", \
      "--timeout", "300", \
      "--keep-alive", "120", \
      "--log-level", "debug", \
+     "--worker-connections", "1000", \
+     "--backlog", "2048", \
      "--bind", "0.0.0.0:7860", \
      "app:app"]

app.py

@@ -45,7 +45,7 @@ class WebSocketBeamStreamer(MultiBeamTextStreamer):
         """Send beam updates through websocket with delay"""
         self.beam_texts[beam_idx] = new_text
         if self.sleep_time > 0:
-            time.sleep(self.sleep_time / 1000)  # Convert milliseconds to seconds
+            eventlet.sleep(self.sleep_time / 1000)  # Convert milliseconds to seconds
         socketio.emit('beam_update', {
             'beam_idx': beam_idx,
             'text': new_text
@@ -65,58 +65,55 @@ def index():
 
 @socketio.on('generate')
 def handle_generation(data):
-    try:
-        app.logger.info("Generation started with data: %s", data)
-        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)
-
-        app.logger.info("Processing with parameters: beams=%d, max_tokens=%d",
-                        num_beams, max_new_tokens)
-
-        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
-        )
+    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)
+            model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
 
-        streamer = WebSocketBeamStreamer(
-            tokenizer=tokenizer,
-            num_beams=num_beams,
-            sleep_time=sleep_time,
-            skip_prompt=True
-        )
-
-        with torch.no_grad():
-            model.generate(
-                **model_inputs,
+            streamer = WebSocketBeamStreamer(
+                tokenizer=tokenizer,
                 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
+                sleep_time=sleep_time,
+                skip_prompt=True
             )
 
-        app.logger.info("Generation completed successfully")
-
-    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')
-
+            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(