more

app.py

@@ -38,14 +38,30 @@ def respond(message, max_new_tokens, temperature, system_message=""):
         ).to("cuda")
 
         # Debug: Inspect input tensors
-        print("[DEBUG] Tokenized input IDs:", inputs["input_ids"].shape)
-        print("[DEBUG] Attention mask:", inputs["attention_mask"].shape)
+        print("[DEBUG] Tokenized inputs:", inputs)
+
+        # Extract `input_ids` and `attention_mask`
+        input_ids = inputs.get("input_ids", None)
+        attention_mask = inputs.get("attention_mask", None)
+
+        # Ensure tensors are present and properly shaped
+        if input_ids is None or attention_mask is None:
+            raise ValueError("Tokenized input is missing `input_ids` or `attention_mask`.")
+        if input_ids.dim() != 2 or attention_mask.dim() != 2:
+            raise ValueError(
+                f"`input_ids` and `attention_mask` must be 2D tensors. "
+                f"Found shapes: input_ids={input_ids.shape}, attention_mask={attention_mask.shape}"
+            )
+
+        # Debug: Shapes of tensors
+        print(f"[DEBUG] input_ids shape: {input_ids.shape}")
+        print(f"[DEBUG] attention_mask shape: {attention_mask.shape}")
 
         # Stream response
         text_streamer = TextStreamer(tokenizer, skip_prompt=True)
         response = model.generate(
-            input_ids=inputs["input_ids"],
-            attention_mask=inputs["attention_mask"],
+            input_ids=input_ids,
+            attention_mask=attention_mask,
             max_new_tokens=max_new_tokens,
             temperature=temperature,
             use_cache=True,