Update app.py

app.py

@@ -46,7 +46,7 @@ def load_model(model_name):
     except Exception as e:
         return f"Erreur lors du chargement du modèle : {str(e)}"
 
-def generate_text(input_text, temperature, top_p, top_k):
+def analyze_next_token(input_text, temperature, top_p, top_k):
     global model, tokenizer
     
     if model is None or tokenizer is None:
@@ -56,39 +56,48 @@ def generate_text(input_text, temperature, top_p, top_k):
     
     try:
         with torch.no_grad():
-            outputs = model.generate(
-                **inputs,
-                max_new_tokens=50,
-                temperature=temperature,
-                top_p=top_p,
-                top_k=top_k,
-                output_attentions=True,
-                return_dict_in_generate=True,
-                output_scores=True
-            )
+            outputs = model(**inputs)
         
-        generated_text = tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
+        last_token_logits = outputs.logits[0, -1, :]
+        probabilities = torch.nn.functional.softmax(last_token_logits, dim=-1)
+        top_k = 5
+        top_probs, top_indices = torch.topk(probabilities, top_k)
+        top_words = [tokenizer.decode([idx.item()]) for idx in top_indices]
+        prob_data = {word: prob.item() for word, prob in zip(top_words, top_probs)}
+        prob_plot = plot_probabilities(prob_data)
         
-        if hasattr(outputs, 'scores') and outputs.scores:
-            last_token_logits = outputs.scores[-1][0]
-            probabilities = torch.nn.functional.softmax(last_token_logits, dim=-1)
-            top_k = 5
-            top_probs, top_indices = torch.topk(probabilities, top_k)
-            top_words = [tokenizer.decode([idx.item()]) for idx in top_indices]
-            prob_data = {word: prob.item() for word, prob in zip(top_words, top_probs)}
-            prob_plot = plot_probabilities(prob_data)
-        else:
-            prob_plot = None
-        
-        if hasattr(outputs, 'attentions') and outputs.attentions:
+        if hasattr(outputs, 'attentions') and outputs.attentions is not None:
             attention_data = torch.mean(torch.stack(outputs.attentions), dim=(0, 1)).cpu().numpy()
             attention_plot = plot_attention(attention_data, tokenizer.convert_ids_to_tokens(inputs["input_ids"][0]))
         else:
             attention_plot = None
         
-        return generated_text, attention_plot, prob_plot
+        return "\n".join([f"{word}: {prob:.4f}" for word, prob in prob_data.items()]), attention_plot, prob_plot
+    except Exception as e:
+        return f"Erreur lors de l'analyse : {str(e)}", None, None
+
+def generate_text(input_text, temperature, top_p, top_k):
+    global model, tokenizer
+    
+    if model is None or tokenizer is None:
+        return "Veuillez d'abord charger un modèle."
+
+    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=512).to(model.device)
+    
+    try:
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=50,
+                temperature=temperature,
+                top_p=top_p,
+                top_k=top_k
+            )
+        
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return generated_text
     except Exception as e:
-        return f"Erreur lors de la génération : {str(e)}", None, None
+        return f"Erreur lors de la génération : {str(e)}"
 
 def plot_attention(attention, tokens):
     fig, ax = plt.subplots(figsize=(10, 10))
@@ -119,10 +128,10 @@ def reset():
     global model, tokenizer
     model = None
     tokenizer = None
-    return "", 1.0, 1.0, 50, None, None, None
+    return "", 1.0, 1.0, 50, None, None, None, None
 
 with gr.Blocks() as demo:
-    gr.Markdown("# Générateur de texte avec visualisation d'attention")
+    gr.Markdown("# Analyse et génération de texte")
     
     with gr.Accordion("Sélection du modèle"):
         model_dropdown = gr.Dropdown(choices=models, label="Choisissez un modèle")
@@ -135,22 +144,28 @@ with gr.Blocks() as demo:
         top_k = gr.Slider(1, 100, value=50, step=1, label="Top-k")
     
     input_text = gr.Textbox(label="Texte d'entrée", lines=3)
-    generate_button = gr.Button("Générer")
+    analyze_button = gr.Button("Analyser le prochain token")
+    generate_button = gr.Button("Générer la suite du texte")
     
-    output_text = gr.Textbox(label="Texte généré", lines=5)
+    next_token_probs = gr.Textbox(label="Probabilités du prochain token")
     
     with gr.Row():
         attention_plot = gr.Plot(label="Visualisation de l'attention")
         prob_plot = gr.Plot(label="Probabilités des tokens suivants")
     
+    generated_text = gr.Textbox(label="Texte généré", lines=5)
+    
     reset_button = gr.Button("Réinitialiser")
     
     load_button.click(load_model, inputs=[model_dropdown], outputs=[load_output])
+    analyze_button.click(analyze_next_token, 
+                         inputs=[input_text, temperature, top_p, top_k], 
+                         outputs=[next_token_probs, attention_plot, prob_plot])
     generate_button.click(generate_text, 
                           inputs=[input_text, temperature, top_p, top_k], 
-                          outputs=[output_text, attention_plot, prob_plot])
+                          outputs=[generated_text])
     reset_button.click(reset, 
-                       outputs=[input_text, temperature, top_p, top_k, output_text, attention_plot, prob_plot])
+                       outputs=[input_text, temperature, top_p, top_k, next_token_probs, attention_plot, prob_plot, generated_text])
 
 if __name__ == "__main__":
     demo.launch()