Update app.py

app.py

@@ -1,7 +1,6 @@
 import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
-import spaces
 import matplotlib.pyplot as plt
 import numpy as np
 from huggingface_hub import login
@@ -25,42 +24,38 @@ tokenizer = None
 def load_model(model_name):
     global model, tokenizer
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="cpu", attn_implementation="eager")
+    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.float16)
     
     # Définir le token de padding s'il n'existe pas
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
         model.config.pad_token_id = model.config.eos_token_id
     
-    return f"Modèle {model_name} chargé avec succès sur CPU."
+    return f"Modèle {model_name} chargé avec succès sur GPU."
 
-@spaces.GPU(duration=300)
 def generate_text(input_text, temperature, top_p, top_k):
     global model, tokenizer
     
-    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=512)
-    input_ids = inputs["input_ids"]
-    attention_mask = inputs["attention_mask"]
+    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=512).to(model.device)
     
     with torch.no_grad():
         outputs = model.generate(
-            input_ids,
-            attention_mask=attention_mask,
+            **inputs,
             max_new_tokens=50,
             temperature=temperature,
             top_p=top_p,
             top_k=top_k,
-            output_attentions=False,
+            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 = model(outputs.sequences[:, -1:]).logits[:, -1, :]
+    last_token_logits = outputs.scores[-1][0]
     
     # Appliquer softmax pour obtenir les probabilités
-    probabilities = torch.nn.functional.softmax(last_token_logits[0], dim=-1)
+    probabilities = torch.nn.functional.softmax(last_token_logits, dim=-1)
     
     # Obtenir les top 5 tokens les plus probables
     top_k = 5
@@ -70,16 +65,18 @@ def generate_text(input_text, temperature, top_p, top_k):
     # 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)}
     
-    # Créer une matrice d'attention factice
+    # Extraire les attentions
+    attentions = torch.cat([att[-1].mean(dim=1) for att in outputs.attentions], dim=0).cpu().numpy()
+    
     attention_data = {
-        'attention': np.random.rand(len(input_ids[0]), len(input_ids[0])).tolist(),
-        'tokens': tokenizer.convert_ids_to_tokens(input_ids[0])
+        'attention': attentions,
+        'tokens': tokenizer.convert_ids_to_tokens(outputs.sequences[0])
     }
     
     return generated_text, plot_attention(attention_data), plot_probabilities(prob_data)
 
 def plot_attention(attention_data):
-    attention = np.array(attention_data['attention'])
+    attention = attention_data['attention']
     tokens = attention_data['tokens']
     
     fig, ax = plt.subplots(figsize=(10, 10))