Update app.py

app.py

@@ -75,6 +75,9 @@ class CassandreChatBot:
         attention_mask = torch.ones_like(input_ids)
 
         try:
+            # Add some debug prints
+            print("Input length:", len(input_ids[0]))
+            
             output = model.generate(
                 input_ids,
                 attention_mask=attention_mask,
@@ -85,11 +88,19 @@ class CassandreChatBot:
                 temperature=temperature,
                 repetition_penalty=repetition_penalty,
                 pad_token_id=tokenizer.pad_token_id,
-                eos_token_id=tokenizer.eos_token_id
+                eos_token_id=tokenizer.eos_token_id,
+                # Add return_dict_in_generate=True to see full output info
+                return_dict_in_generate=True,
+                output_scores=True
             )
-
-            # Just decode and display the generated text directly
-            generated_text = tokenizer.decode(output[0])
+            
+            # Print debug info about output
+            print("Output sequence length:", len(output.sequences[0]))
+            print("New tokens generated:", len(output.sequences[0]) - len(input_ids[0]))
+            
+            # Try decoding only the new tokens
+            generated_text = tokenizer.decode(output.sequences[0][len(input_ids[0]):])
+            
             generated_text = '<h2 style="text-align:center">Réponse</h3>\n<div class="generation">' + format_references(generated_text) + "</div>"
             fiches_html = '<h2 style="text-align:center">Sources</h3>\n' + fiches_html
             return generated_text, fiches_html