Upload app.py

app.py

@@ -107,7 +107,7 @@ def generate(model, tokenizer, prompt, steps=128, max_new_tokens=128, block_size
     if isinstance(prompt, torch.Tensor):
         x = prompt.to(device).long()
     else:
-        if isinstance(prompt, (list, tuple)):
+        if isinstance(prompt[0], (list, tuple)):
             max_len = max(len(p) for p in prompt)
             x = torch.full((len(prompt), max_len), pad_id, device=device, dtype=torch.long)
             for i, p in enumerate(prompt):
@@ -165,11 +165,12 @@ def generate(model, tokenizer, prompt, steps=128, max_new_tokens=128, block_size
             if capture_interval > 0 and total_step % capture_interval == 0:
                 intermediates.append(x.clone())
             total_step += 1
-        if tokenizer.eos_token_id is not None:
-            finished |= (x_blk_new == tokenizer.eos_token_id).any(dim=1)
+            if tokenizer.eos_token_id is not None:
+                finished |= (x_blk_new == tokenizer.eos_token_id).any(dim=1)
+        generated += cur_len
         if finished.all():
             break
-        generated += cur_len
+
     if capture_interval > 0:
         return x, intermediates
     return x
@@ -184,7 +185,7 @@ def generate_stream(model, tokenizer, prompt, steps=128, max_new_tokens=128, blo
     if isinstance(prompt, torch.Tensor):
         x = prompt.to(device).long()
     else:
-        if isinstance(prompt, (list, tuple)):
+        if isinstance(prompt[0], (list, tuple)):
             max_len = max(len(p) for p in prompt)
             x = torch.full((len(prompt), max_len), pad_id, device=device, dtype=torch.long)
             for i, p in enumerate(prompt):
@@ -261,7 +262,7 @@ def generate_stream(model, tokenizer, prompt, steps=128, max_new_tokens=128, blo
                 finished |= (x_blk_new == tokenizer.eos_token_id).any(dim=1)
             if finished.all():
                 break
-            generated += cur_len
+        generated += cur_len
             
         if finished.all():
             break
@@ -370,7 +371,7 @@ def generate_text_stream():
         return jsonify({"error": "Model workspace offline"}), 503
         
     data = request.get_json() or {}
-    if 'prompt' not in data:
+    if not data or 'prompt' not in data:
         return jsonify({"error": "Missing 'prompt' operational field"}), 400
         
     prompt = data['prompt']
@@ -422,7 +423,7 @@ def generate_text_sse():
         return jsonify({"error": "Model workspace offline"}), 503
         
     data = request.get_json() or {}
-    if 'prompt' not in data:
+    if not data or 'prompt' not in data:
         return jsonify({"error": "Missing 'prompt' operational field"}), 400
         
     prompt = data['prompt']