Update app.py

app.py

@@ -7,30 +7,12 @@ import matplotlib.pyplot as plt
 import seaborn as sns
 import numpy as np
 import time
-from langdetect import detect
 
 # Authentification
 login(token=os.environ["HF_TOKEN"])
 
-# Liste des modèles
-models = [
-    "meta-llama/Llama-2-13b-hf",
-    "meta-llama/Llama-2-7b-hf",
-    "meta-llama/Llama-2-70b-hf",
-    "meta-llama/Meta-Llama-3-8B",
-    "meta-llama/Llama-3.2-3B",
-    "meta-llama/Llama-3.1-8B",
-    "mistralai/Mistral-7B-v0.1",
-    "mistralai/Mixtral-8x7B-v0.1",
-    "mistralai/Mistral-7B-v0.3",
-    "google/gemma-2-2b",
-    "google/gemma-2-9b",
-    "google/gemma-2-27b",
-    "croissantllm/CroissantLLMBase"
-]
-
-# Dictionnaire des langues supportées par modèle
-model_languages = {
+# Liste des modèles et leurs langues supportées
+models_and_languages = {
     "meta-llama/Llama-2-13b-hf": ["en"],
     "meta-llama/Llama-2-7b-hf": ["en"],
     "meta-llama/Llama-2-70b-hf": ["en"],
@@ -49,6 +31,7 @@ model_languages = {
 # Variables globales
 model = None
 tokenizer = None
+selected_language = None
 
 def load_model(model_name, progress=gr.Progress()):
     global model, tokenizer
@@ -57,21 +40,40 @@ def load_model(model_name, progress=gr.Progress()):
         tokenizer = AutoTokenizer.from_pretrained(model_name)
         progress(0.5, desc="Chargement du modèle")
         
-        # Configuration générique pour tous les modèles
-        model = AutoModelForCausalLM.from_pretrained(
-            model_name,
-            torch_dtype=torch.float16,
-            device_map="auto",
-            low_cpu_mem_usage=True
-        )
+        # Configurations spécifiques par modèle
+        if "mixtral" in model_name.lower():
+            model = AutoModelForCausalLM.from_pretrained(
+                model_name,
+                torch_dtype=torch.float16,
+                device_map="auto",
+                load_in_8bit=True
+            )
+        elif "llama" in model_name.lower() or "mistral" in model_name.lower():
+            model = AutoModelForCausalLM.from_pretrained(
+                model_name,
+                torch_dtype=torch.float16,
+                device_map="auto"
+            )
+        else:
+            model = AutoModelForCausalLM.from_pretrained(
+                model_name,
+                torch_dtype=torch.float16,
+                device_map="auto"
+            )
         
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
         
         progress(1.0, desc="Modèle chargé")
-        return f"Modèle {model_name} chargé avec succès."
+        available_languages = models_and_languages[model_name]
+        return f"Modèle {model_name} chargé avec succès. Langues disponibles : {', '.join(available_languages)}", gr.Dropdown.update(choices=available_languages, value=available_languages[0], visible=True)
     except Exception as e:
-        return f"Erreur lors du chargement du modèle : {str(e)}"
+        return f"Erreur lors du chargement du modèle : {str(e)}", gr.Dropdown.update(visible=False)
+
+def set_language(lang):
+    global selected_language
+    selected_language = lang
+    return f"Langue sélectionnée : {lang}"
 
 def ensure_token_display(token):
     """Assure que le token est affiché correctement."""
@@ -80,29 +82,23 @@ def ensure_token_display(token):
     return token
 
 def analyze_next_token(input_text, temperature, top_p, top_k):
-    global model, tokenizer
+    global model, tokenizer, selected_language
     
     if model is None or tokenizer is None:
         return "Veuillez d'abord charger un modèle.", None, None
 
-    # Détection de la langue
-    detected_lang = detect(input_text)
-    if detected_lang not in model_languages.get(model.config._name_or_path, []):
-        return f"Langue détectée ({detected_lang}) non supportée par ce modèle.", None, None
-
-    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=512).to(model.device)
+    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=512)
     
     try:
         with torch.no_grad():
             outputs = model(**inputs)
         
         last_token_logits = outputs.logits[0, -1, :]
-        probabilities = torch.nn.functional.softmax(last_token_logits / temperature, dim=-1)
+        probabilities = torch.nn.functional.softmax(last_token_logits, dim=-1)
         
-        top_k = min(top_k, probabilities.size(-1))
+        top_k = 10
         top_probs, top_indices = torch.topk(probabilities, top_k)
         top_words = [ensure_token_display(tokenizer.decode([idx.item()])) for idx in top_indices]
-        
         prob_data = {word: prob.item() for word, prob in zip(top_words, top_probs)}
         
         prob_text = "Prochains tokens les plus probables :\n\n"
@@ -117,27 +113,22 @@ def analyze_next_token(input_text, temperature, top_p, top_k):
         return f"Erreur lors de l'analyse : {str(e)}", None, None
 
 def generate_text(input_text, temperature, top_p, top_k):
-    global model, tokenizer
+    global model, tokenizer, selected_language
     
     if model is None or tokenizer is None:
         return "Veuillez d'abord charger un modèle."
 
-    # Détection de la langue
-    detected_lang = detect(input_text)
-    if detected_lang not in model_languages.get(model.config._name_or_path, []):
-        return f"Langue détectée ({detected_lang}) non supportée par ce modèle."
-
-    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=512).to(model.device)
+    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=512)
     
     try:
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=50,
-            do_sample=True,
-            temperature=temperature,
-            top_p=top_p,
-            top_k=top_k
-        )
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=1,
+                temperature=temperature,
+                top_p=top_p,
+                top_k=top_k
+            )
         
         generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
         return generated_text
@@ -172,7 +163,7 @@ def plot_attention(input_ids, last_token_logits):
     top_attention_scores, _ = torch.topk(attention_scores, top_k)
     
     fig, ax = plt.subplots(figsize=(14, 7))
-    sns.heatmap(top_attention_scores.unsqueeze(0).cpu().numpy(), annot=True, cmap="YlOrRd", cbar=True, ax=ax, fmt='.2%')
+    sns.heatmap(top_attention_scores.unsqueeze(0).numpy(), annot=True, cmap="YlOrRd", cbar=True, ax=ax, fmt='.2%')
     ax.set_xticklabels(input_tokens[-top_k:], rotation=45, ha="right", fontsize=10)
     ax.set_yticklabels(["Attention"], rotation=0, fontsize=10)
     ax.set_title("Scores d'attention pour les derniers tokens", fontsize=16)
@@ -185,18 +176,21 @@ def plot_attention(input_ids, last_token_logits):
     return fig
 
 def reset():
-    global model, tokenizer
+    global model, tokenizer, selected_language
     model = None
     tokenizer = None
-    return "", 1.0, 1.0, 50, None, None, None, None
+    selected_language = None
+    return "", 1.0, 1.0, 50, None, None, None, None, gr.Dropdown.update(visible=False), ""
 
 with gr.Blocks() as demo:
-    gr.Markdown("# Analyse et génération de texte avec LLM")
+    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")
+        model_dropdown = gr.Dropdown(choices=list(models_and_languages.keys()), label="Choisissez un modèle")
         load_button = gr.Button("Charger le modèle")
         load_output = gr.Textbox(label="Statut du chargement")
+        language_dropdown = gr.Dropdown(label="Choisissez une langue", visible=False)
+        language_output = gr.Textbox(label="Langue sélectionnée")
     
     with gr.Row():
         temperature = gr.Slider(0.1, 2.0, value=1.0, label="Température")
@@ -212,12 +206,13 @@ with gr.Blocks() as demo:
         attention_plot = gr.Plot(label="Visualisation de l'attention")
         prob_plot = gr.Plot(label="Probabilités des tokens suivants")
     
-    generate_button = gr.Button("Générer la suite du texte")
+    generate_button = gr.Button("Générer le prochain mot")
     generated_text = gr.Textbox(label="Texte généré")
     
     reset_button = gr.Button("Réinitialiser")
     
-    load_button.click(load_model, inputs=[model_dropdown], outputs=[load_output])
+    load_button.click(load_model, inputs=[model_dropdown], outputs=[load_output, language_dropdown])
+    language_dropdown.change(set_language, inputs=[language_dropdown], outputs=[language_output])
     analyze_button.click(analyze_next_token, 
                          inputs=[input_text, temperature, top_p, top_k], 
                          outputs=[next_token_probs, attention_plot, prob_plot])
@@ -225,7 +220,7 @@ with gr.Blocks() as demo:
                           inputs=[input_text, temperature, top_p, top_k], 
                           outputs=[generated_text])
     reset_button.click(reset, 
-                       outputs=[input_text, temperature, top_p, top_k, next_token_probs, attention_plot, prob_plot, generated_text])
+                       outputs=[input_text, temperature, top_p, top_k, next_token_probs, attention_plot, prob_plot, generated_text, language_dropdown, language_output])
 
 if __name__ == "__main__":
     demo.launch()