more

app.py

@@ -30,38 +30,24 @@ def respond(message, max_new_tokens, temperature, system_message=""):
         print("[DEBUG] Messages:", messages)
 
         # Tokenize inputs
-        inputs = tokenizer.apply_chat_template(
+        input_ids = tokenizer.apply_chat_template(
             messages,
             tokenize=True,
             add_generation_prompt=True,
             return_tensors="pt",
         ).to("cuda")
 
-        # Debug: Inspect input tensors
-        print("[DEBUG] Tokenized inputs:", inputs)
+        # Debug: Inspect input tensor
+        print("[DEBUG] input_ids:", input_ids)
 
-        # 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}")
+        # Ensure the input tensor has the correct dimensions
+        if input_ids.dim() != 2:
+            raise ValueError(f"`input_ids` must be a 2D tensor. Found shape: {input_ids.shape}")
 
         # Stream response
         text_streamer = TextStreamer(tokenizer, skip_prompt=True)
         response = model.generate(
             input_ids=input_ids,
-            attention_mask=attention_mask,
             max_new_tokens=max_new_tokens,
             temperature=temperature,
             use_cache=True,
@@ -74,6 +60,7 @@ def respond(message, max_new_tokens, temperature, system_message=""):
         print("[ERROR]", str(e))
         return f"Error: {str(e)}"
 
+
 # Gradio UI
 demo = gr.Interface(
     fn=respond,