more

app.py

@@ -50,15 +50,20 @@ def respond(
             tokenize=True,
             add_generation_prompt=True,
             return_tensors="pt",
-        ).to("cuda")
+        )
+
+        # Ensure tensor shapes are correct
+        input_ids = inputs["input_ids"].squeeze(0).to("cuda")
+        attention_mask = inputs["attention_mask"].squeeze(0).to("cuda")
 
         # Debug: Print tokenized inputs
-        print("[DEBUG] Tokenized inputs:", inputs)
+        print("[DEBUG] Tokenized input_ids shape:", input_ids.shape)
+        print("[DEBUG] Tokenized attention_mask shape:", attention_mask.shape)
 
         # Generate response
         output_ids = model.generate(
-            input_ids=inputs["input_ids"],
-            attention_mask=inputs["attention_mask"],
+            input_ids=input_ids,
+            attention_mask=attention_mask,
             max_new_tokens=max_tokens,
             temperature=temperature,
             top_p=top_p,