GPT2

app.py

@@ -2,8 +2,7 @@ import eventlet
 eventlet.monkey_patch(socket=True, select=True, thread=True)
 
 import eventlet.wsgi
-
-from flask import Flask, render_template
+from flask import Flask, render_template, request
 from flask_socketio import SocketIO
 from transformers import MultiBeamTextStreamer, AutoTokenizer, AutoModelForCausalLM
 import torch
@@ -12,23 +11,44 @@ app = Flask(__name__)
 socketio = SocketIO(
     app,
     async_mode='eventlet',
-    message_queue=None,  # Explicitly set to None for single-worker setup
+    message_queue=None,
     ping_timeout=60,
     ping_interval=25,
     cors_allowed_origins="*",
     logger=True,
     engineio_logger=True,
-    async_handlers=True  # Enable async handlers
-)
-# Initialize model and tokenizer
-MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    torch_dtype="auto",
-    device_map="auto"
+    async_handlers=True
 )
 
+# Initialize models and tokenizers
+MODELS = {
+    "qwen": {
+        "name": "Qwen/Qwen2.5-0.5B-Instruct",
+        "tokenizer": None,
+        "model": None,
+        "uses_chat_template": True  # Qwen uses chat template
+    },
+    "gpt2": {
+        "name": "gpt2",
+        "tokenizer": None,
+        "model": None,
+        "uses_chat_template": False  # GPT2 doesn't use chat template
+    }
+}
+
+# Load models and tokenizers
+for model_key, model_info in MODELS.items():
+    model_info["tokenizer"] = AutoTokenizer.from_pretrained(model_info["name"])
+    model_info["model"] = AutoModelForCausalLM.from_pretrained(
+        model_info["name"],
+        torch_dtype="auto",
+        device_map="auto"
+    )
+    # Add pad token for GPT2 if it doesn't have one
+    if model_key == "gpt2" and model_info["tokenizer"].pad_token is None:
+        model_info["tokenizer"].pad_token = model_info["tokenizer"].eos_token
+        model_info["model"].config.pad_token_id = model_info["model"].config.eos_token_id
+
 
 class WebSocketBeamStreamer(MultiBeamTextStreamer):
     """Custom streamer that sends updates through websockets with adjustable speed"""
@@ -42,22 +62,19 @@ class WebSocketBeamStreamer(MultiBeamTextStreamer):
             on_beam_finished=self.on_beam_finished
         )
         self.beam_texts = {i: "" for i in range(num_beams)}
-        self.sleep_time = sleep_time  # Sleep time in milliseconds
+        self.sleep_time = sleep_time
 
     def on_beam_update(self, beam_idx: int, new_text: str):
-        """Send beam updates through websocket with delay"""
         self.beam_texts[beam_idx] = new_text
         if self.sleep_time > 0:
-            eventlet.sleep(self.sleep_time / 1000)  # Convert milliseconds to seconds
-        # Force immediate emit and flush
+            eventlet.sleep(self.sleep_time / 1000)
         socketio.emit('beam_update', {
             'beam_idx': beam_idx,
             'text': new_text
         }, namespace='/', callback=lambda: eventlet.sleep(0))
-        socketio.sleep(0)  # Force context switch
+        socketio.sleep(0)
 
     def on_beam_finished(self, final_text: str):
-        """Send completion notification through websocket"""
         socketio.emit('beam_finished', {
             'text': final_text
         })
@@ -70,31 +87,39 @@ def index():
 
 @socketio.on('generate')
 def handle_generation(data):
-    # Emit a generation start event
     socketio.emit('generation_started')
 
     prompt = data['prompt']
+    model_name = data.get('model', 'qwen')  # Default to qwen if not specified
     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
-    )
+    sleep_time = data.get('sleep_time', 0)
+
+    # Get the selected model info
+    model_info = MODELS[model_name]
+    model = model_info["model"]
+    tokenizer = model_info["tokenizer"]
+
+    # Prepare input text based on model type
+    if model_info["uses_chat_template"]:
+        # For Qwen, use chat template
+        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
+        )
+    else:
+        # For GPT2, use direct prompt
+        text = prompt
 
     # Prepare inputs
     model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
 
-    # Initialize streamer with sleep time
+    # Initialize streamer
     streamer = WebSocketBeamStreamer(
         tokenizer=tokenizer,
         num_beams=num_beams,
@@ -113,10 +138,11 @@ def handle_generation(data):
                 output_scores=True,
                 return_dict_in_generate=True,
                 early_stopping=True,
-                streamer=streamer
+                streamer=streamer,
+                pad_token_id=tokenizer.pad_token_id,
+                eos_token_id=tokenizer.eos_token_id
             )
     except Exception as e:
         socketio.emit('generation_error', {'error': str(e)})
     finally:
-        # Emit generation completed event
-        socketio.emit('generation_completed')
+        socketio.emit('generation_completed')

templates/index.html

@@ -357,6 +357,14 @@
                     <label for="max_tokens">Max tokens</label>
                     <input type="number" id="max_tokens" value="512" min="1">
                 </div>
+
+                <div class="input-group">
+                    <label for="model_select">Model</label>
+                    <select id="model_select" class="form-select">
+                        <option value="gpt2">GPT-2</option>
+                        <option value="qwen">Qwen</option>
+                    </select>
+                </div>
             </div>
 
             <div class="slider-container">
@@ -517,6 +525,7 @@
             resetConnection();
 
             const prompt = document.getElementById('prompt').value;
+            const model = document.getElementById('model_select').value;
             const numBeams = parseInt(document.getElementById('num_beams').value);
             const maxTokens = parseInt(document.getElementById('max_tokens').value);
             const sleepTime = parseInt(document.getElementById('sleep_time').value);
@@ -527,6 +536,7 @@
 
             socket.emit('generate', {
                 prompt: prompt,
+                model: model,
                 num_beams: numBeams,
                 max_tokens: maxTokens,
                 sleep_time: sleepTime