Update app.py

app.py

@@ -75,9 +75,6 @@ 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,
@@ -88,18 +85,11 @@ class CassandreChatBot:
                 temperature=temperature,
                 repetition_penalty=repetition_penalty,
                 pad_token_id=tokenizer.pad_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
+                eos_token_id=tokenizer.eos_token_id
             )
-            
-            # 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]):])
+
+            # Only decode the new tokens by slicing from the input length
+            generated_text = tokenizer.decode(output[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