Update app.py

app.py

@@ -6,7 +6,9 @@ import matplotlib.pyplot as plt
 import numpy as np
 from huggingface_hub import login
 import os
+
 login(token=os.environ["HF_TOKEN"])
+
 # Liste des modèles
 models = [
     "meta-llama/Llama-2-13b", "meta-llama/Llama-2-7b", "meta-llama/Llama-2-70b",
@@ -23,14 +25,13 @@ tokenizer = None
 def load_model(model_name):
     global model, tokenizer
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name)
-    return f"Modèle {model_name} chargé avec succès."
+    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="cpu")
+    return f"Modèle {model_name} chargé avec succès sur CPU."
 
-@spaces.GPU(duration=300)
 def generate_text(input_text, temperature, top_p, top_k):
     global model, tokenizer
     
-    inputs = tokenizer(input_text, return_tensors="pt").to("cuda")
+    inputs = tokenizer(input_text, return_tensors="pt")
     
     with torch.no_grad():
         outputs = model.generate(
@@ -46,8 +47,8 @@ def generate_text(input_text, temperature, top_p, top_k):
     generated_text = tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
     
     # Extraire les attentions et les logits
-    attentions = outputs.attentions[-1][0][-1].cpu().numpy()
-    logits = outputs.scores[-1][0].cpu()
+    attentions = outputs.attentions[-1][0][-1].numpy()
+    logits = outputs.scores[-1][0]
     
     # Visualiser l'attention
     plt.figure(figsize=(10, 10))