Update app.py

app.py

@@ -32,48 +32,68 @@ tokenizer = None
 
 def load_model(model_name):
     global model, tokenizer
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="auto")
-    if tokenizer.pad_token is None:
-        tokenizer.pad_token = tokenizer.eos_token
-    return f"Modèle {model_name} chargé avec succès."
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="auto", attn_implementation="eager")
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+        return f"Modèle {model_name} chargé avec succès."
+    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):
     global model, tokenizer
     
+    if model is None or tokenizer is None:
+        return "Veuillez d'abord charger un modèle.", None, None
+
     inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=512).to(model.device)
     
-    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
-        )
-    
-    generated_text = tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
-    
-    # Obtenir les logits pour le dernier token généré
-    last_token_logits = outputs.scores[-1][0]
-    
-    # Appliquer softmax pour obtenir les probabilités
-    probabilities = torch.nn.functional.softmax(last_token_logits, dim=-1)
-    
-    # Obtenir les top 5 tokens les plus probables
-    top_k = 5
-    top_probs, top_indices = torch.topk(probabilities, top_k)
-    top_words = [tokenizer.decode([idx.item()]) for idx in top_indices]
-    
-    # Préparer les données pour le graphique des probabilités
-    prob_data = {word: prob.item() for word, prob in zip(top_words, top_probs)}
-    
-    # Extraire les attentions (moyenne sur toutes les couches et têtes d'attention)
-    attentions = torch.mean(torch.stack(outputs.attentions), dim=(0, 1)).cpu().numpy()
-    
-    return generated_text, plot_attention(attentions, tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])), plot_probabilities(prob_data)
+    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
+            )
+        
+        generated_text = tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
+        
+        # Obtenir les logits pour le dernier token généré
+        if outputs.scores:
+            last_token_logits = outputs.scores[-1][0]
+            
+            # Appliquer softmax pour obtenir les probabilités
+            probabilities = torch.nn.functional.softmax(last_token_logits, dim=-1)
+            
+            # Obtenir les top 5 tokens les plus probables
+            top_k = 5
+            top_probs, top_indices = torch.topk(probabilities, top_k)
+            top_words = [tokenizer.decode([idx.item()]) for idx in top_indices]
+            
+            # Préparer les données pour le graphique des probabilités
+            prob_data = {word: prob.item() for word, prob in zip(top_words, top_probs)}
+            
+            # Extraire les attentions (moyenne sur toutes les couches et têtes d'attention)
+            if outputs.attentions:
+                attentions = torch.mean(torch.stack(outputs.attentions), dim=(0, 1)).cpu().numpy()
+                attention_plot = plot_attention(attentions, tokenizer.convert_ids_to_tokens(inputs["input_ids"][0]))
+            else:
+                attention_plot = None
+            
+            prob_plot = plot_probabilities(prob_data)
+        else:
+            attention_plot = None
+            prob_plot = None
+        
+        return generated_text, attention_plot, prob_plot
+    except Exception as e:
+        return f"Erreur lors de la génération : {str(e)}", None, None
 
 def plot_attention(attention, tokens):
     fig, ax = plt.subplots(figsize=(10, 10))
@@ -101,6 +121,9 @@ def plot_probabilities(prob_data):
     return fig
 
 def reset():
+    global model, tokenizer
+    model = None
+    tokenizer = None
     return "", 1.0, 1.0, 50, None, None, None
 
 with gr.Blocks() as demo: