Update app.py

app.py

@@ -4,10 +4,14 @@ import pandas as pd
 import json
 from transformers import AutoModel, AutoProcessor
 
+# Vérifier si CUDA est disponible (GPU)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"🔹 Utilisation du périphérique : {device}")
+
 # Charger le modèle Marqo avec gestion d'erreurs
 model_name = "Marqo/marqo-ecommerce-embeddings-L"
 try:
-    model = AutoModel.from_pretrained(model_name, trust_remote_code=True)
+    model = AutoModel.from_pretrained(model_name, trust_remote_code=True).to(device)
     processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
 except Exception as e:
     print(f"❌ Erreur lors du chargement du modèle : {e}")
@@ -19,7 +23,7 @@ def load_products_from_json():
         with open("products.json", "r", encoding="utf-8") as f:
             data = json.load(f)
         
-        products = data.get("products", [])
+        products = data.get("products", [])[:1000]  # Limiter à 1000 produits max pour éviter crash mémoire
         structured_products = []
         
         for product in products:
@@ -28,7 +32,7 @@ def load_products_from_json():
                 structured_products.append({
                     "id": product.get("id", "N/A"),
                     "title": product.get("title", "Produit inconnu"),
-                    "description": product.get("description", "")[:200],
+                    "description": product.get("description", "")[:100],
                     "category": product.get("product_type", "Inconnu"),
                     "brand": product.get("brand", "Sans marque"),
                     "price": price,
@@ -50,7 +54,7 @@ def load_products_from_json():
 # Charger les produits
 products_data = load_products_from_json()
 
-def get_text_embeddings(texts, batch_size=8):
+def get_text_embeddings(texts, batch_size=16):
     if model is None or processor is None:
         print("❌ Modèle non chargé, impossible de générer les embeddings.")
         return torch.empty((0, 1024))
@@ -58,10 +62,10 @@ def get_text_embeddings(texts, batch_size=8):
     embeddings = []
     for i in range(0, len(texts), batch_size):
         batch = texts[i:i+batch_size]
-        processed = processor(text=batch, return_tensors="pt", truncation=True, max_length=64, padding=True)
+        processed = processor(text=batch, return_tensors="pt", truncation=True, max_length=64, padding=True).to(device)
         with torch.no_grad():
             batch_embeddings = model.get_text_features(processed["input_ids"], normalize=True)
-        embeddings.extend(batch_embeddings)
+        embeddings.extend(batch_embeddings.cpu())  # Stocker en CPU pour libérer la mémoire GPU
     return torch.stack(embeddings)
 
 if not products_data:
@@ -71,8 +75,6 @@ if not products_data:
 else:
     products_df = pd.DataFrame(products_data)
     title_embeddings = get_text_embeddings([prod["title"] for prod in products_data])
-    description_embeddings = get_text_embeddings([prod["description"] for prod in products_data])
-    category_embeddings = get_text_embeddings([prod["category"] for prod in products_data])
 
 def search_products(query, category, min_price, max_price):
     print(f"🔎 Recherche déclenchée avec: {query}, Catégorie: {category}, Prix: {min_price}-{max_price}")
@@ -81,20 +83,12 @@ def search_products(query, category, min_price, max_price):
         return pd.DataFrame()
     
     query_embedding = get_text_embeddings([query])[0]
-    
-    # Calcul des similarités pondérées
-    title_sim = torch.nn.functional.cosine_similarity(query_embedding, title_embeddings, dim=1) * 0.5
-    desc_sim = torch.nn.functional.cosine_similarity(query_embedding, description_embeddings, dim=1) * 0.3
-    cat_sim = torch.nn.functional.cosine_similarity(query_embedding, category_embeddings, dim=1) * 0.2
-    total_sim = title_sim + desc_sim + cat_sim
-    
-    # Normalisation des scores
-    normalized_similarities = (total_sim - total_sim.min()) / (total_sim.max() - total_sim.min())
+    title_sim = torch.nn.functional.cosine_similarity(query_embedding, title_embeddings, dim=1)
+    normalized_similarities = (title_sim - title_sim.min()) / (title_sim.max() - title_sim.min())
     
     results = products_df.copy()
     results["score"] = normalized_similarities.cpu().numpy()
     
-    # Filtrage strict des résultats
     filtered_results = results[
         (results["price"] >= min_price) &
         (results["price"] <= max_price) &
@@ -130,5 +124,3 @@ if model is not None and not products_df.empty:
     app = create_ui()
     app.queue()
     app.launch(server_name="0.0.0.0", server_port=7860, share=True)
-else:
-    print("❌ L'application n'a pas pu être initialisée, vérifiez les erreurs ci-dessus.")