Update fashion_query.py

fashion_query.py

@@ -1,10 +1,11 @@
-import os
 import pandas as pd
 import random
+import os
 from transformers import DistilBertTokenizer, DistilBertModel
 from huggingface_hub import login
+import torch
 
-# Check if the token is available
+# Load environment variables
 HUGGINGFACE_TOKEN = os.getenv('HUGGINGFACE_TOKEN')
 if not HUGGINGFACE_TOKEN:
     raise EnvironmentError("HUGGINGFACE_TOKEN environment variable is not set.")
@@ -17,7 +18,16 @@ model_name = "distilbert-base-uncased"
 tokenizer = DistilBertTokenizer.from_pretrained(model_name, use_auth_token=HUGGINGFACE_TOKEN)
 model = DistilBertModel.from_pretrained(model_name, use_auth_token=HUGGINGFACE_TOKEN)
 
-# Load the fashion dataset
+def predict_class(query, system_message):
+    """Predict the class of the query based on the provided system message."""
+    inputs = tokenizer(system_message + " " + query, return_tensors='pt')
+    with torch.no_grad():
+        outputs = model(**inputs)
+    logits = outputs.last_hidden_state[:, 0, :]  # Using [CLS] token's embedding
+    probabilities = torch.nn.functional.softmax(logits, dim=-1)
+    predicted_class = torch.argmax(probabilities).item()
+    return predicted_class
+
 def load_fashion_dataset():
     try:
         fashion_df = pd.read_csv('fashion.csv')
@@ -31,41 +41,28 @@ def load_fashion_dataset():
     except Exception as e:
         raise Exception(f"Error loading fashion dataset: {e}")
 
-# Define different response styles
-fashion_response_templates = [
-    lambda row: (
-        f"The {row['ProductName']} is perfect for a stylish {row['Category']}. "
-        f"Available for ${row['Price']}, it's known for its {row['Description']}. "
-        f"Would you like more details on the {row['ProductName']}?"
-    ),
-    lambda row: (
-        f"Discover the {row['ProductName']}! Priced at ${row['Price']} and famous for its {row['Description']}. "
-        f"What are your thoughts on this {row['Category']}?"
-    ),
-    # Add more templates as needed
-]
-
-# Generate a dynamic response
 def generate_fashion_response(row):
-    template = random.choice(fashion_response_templates)
+    templates = [
+        lambda r: (f"The {r['ProductName']} is perfect for a stylish {r['Category']}. "
+                   f"Available for ${r['Price']}, it's known for its {r['Description']}. "
+                   f"Would you like more details on the {r['ProductName']}?"),
+        lambda r: (f"Discover the {r['ProductName']}! Priced at ${r['Price']} and famous for its {r['Description']}. "
+                   f"What are your thoughts on this {r['Category']}?"),
+    ]
+    template = random.choice(templates)
     return template(row)
 
-# Extract filters from the query
 def extract_fashion_filters(query):
     filters = {}
     query_lower = query.lower()
-
     if 'best' in query_lower and 'rating' in query_lower:
         filters['Rating'] = 'max'
-
     if 'dresses' in query_lower:
         filters['Category'] = 'dress'
     elif 'shoes' in query_lower:
         filters['Category'] = 'shoes'
-    
     return filters
 
-# Apply filters to the fashion DataFrame
 def apply_fashion_filters(df, filters):
     for key, value in filters.items():
         if key == 'Rating' and value == 'max':
@@ -74,30 +71,35 @@ def apply_fashion_filters(df, filters):
             df = df[df[key].str.contains(value, case=False, na=False)]
     return df
 
-# Query fashion based on user input
 def query_fashion(user_query, n_results=5):
     fashion_df = load_fashion_dataset()
     filtered_df = apply_fashion_filters(fashion_df, extract_fashion_filters(user_query))
-    
-    # Check if 'Rating' column exists before sorting
     if 'Rating' in filtered_df.columns:
         sorted_df = filtered_df.sort_values(by='Rating', ascending=False)
     else:
-        sorted_df = filtered_df  # Skip sorting if 'Rating' is not available
-
-    # Return the top N results
+        sorted_df = filtered_df
     return sorted_df.head(n_results)
 
-# Function to integrate the chatbot responses
-def fashion_chatbot_response(user_query):
-    # Query the fashion data
-    fashion_results = query_fashion(user_query)
+def fashion_agent_response(user_query):
+    system_message = """You are a knowledgeable fashion agent. Your responsibilities include:
+    1. Handling all fashion-related queries.
+    2. Providing information about clothing, accessories, and trends.
+    3. Assisting with product details, styles, and prices.
+    4. Offering fashion recommendations based on user preferences."""
+
+    predicted_class = predict_class(user_query, system_message)
+    responses = {
+        0: "Information about fashion trends.",
+        1: "Details about clothing styles and prices.",
+        2: "Assistance with fashion recommendations."
+    }
 
+    response = responses.get(predicted_class, "I am not sure how to help with that.")
+    fashion_results = query_fashion(user_query)
     if not fashion_results.empty:
-        response = ""
         for _, row in fashion_results.iterrows():
-            response += generate_fashion_response(row) + "\n"
-        return response
+            response += "\n" + generate_fashion_response(row)
+    else:
+        response += "\nSorry, I couldn't find any fashion items matching your query."
 
-    # If no results, generate a fallback response
-    return "Sorry, I couldn't find any fashion items matching your query."
+    return response