Update app.py

app.py

@@ -3,46 +3,54 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import requests
 
-# Load DeepSeek-R1 model and tokenizer
-model_name = "deepseek-ai/DeepSeek-R1"
+# Load DeepSeek R1 model with Accelerate enabled
+model_name = "deepseek-ai/deepseek-moe-8b-chat"
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name, 
+    torch_dtype=torch.float16, 
+    device_map="auto", 
+    trust_remote_code=True, 
+    low_cpu_mem_usage=True  # Helps reduce memory consumption
+)
+
+# 🎯 Extract interests using DeepSeek R1
+def extract_interests(text):
+    prompt = f"Extract key interests from this request: '{text}'. Give 3-5 relevant words only."
+
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")  # Runs on GPU if available
+    outputs = model.generate(**inputs, max_length=100)
+    interests = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-# 🎁 Web search function to get gift suggestions from various e-commerce sites
-def search_gifts(query):
-    search_query = query.replace(" ", "+")
+    return interests.split(", ")  # Return as a list of keywords
+
+# 🎁 Search for gift suggestions from various e-commerce sites
+def search_gifts(interests):
+    query = "+".join(interests)
     return {
-        "Amazon": f"https://www.amazon.in/s?k={search_query}",
-        "Flipkart": f"https://www.flipkart.com/search?q={search_query}",
-        "IGP": f"https://www.igp.com/search?q={search_query}",
-        "IndiaMart": f"https://dir.indiamart.com/search.mp?ss={search_query}",
+        "Amazon": f"https://www.amazon.in/s?k={query}",
+        "Flipkart": f"https://www.flipkart.com/search?q={query}",
+        "IGP": f"https://www.igp.com/search?q={query}",
+        "IndiaMart": f"https://dir.indiamart.com/search.mp?ss={query}"
     }
 
-# 🎯 Generate gift recommendations using AI
+# 🎯 Main function for generating gift recommendations
 def recommend_gifts(text):
     if not text:
         return "Please enter a description."
 
-    # Prepare input prompt for the model
-    prompt = f"Suggest the best gifts for: '{text}'"
-
-    # Generate response using the model
-    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
-    outputs = model.generate(inputs.input_ids, max_length=200, do_sample=True)
-    recommendation = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-    # Search for gifts on shopping websites
-    product_links = search_gifts(recommendation)
+    interests = extract_interests(text)  # Extract relevant interests
+    links = search_gifts(interests)  # Generate shopping links
 
     return {
-        "Predicted Gift": recommendation,
-        "Gift Suggestions": product_links
+        "Predicted Interests": interests,
+        "Gift Suggestions": links
     }
 
-# 🎨 Gradio UI for easy interaction
+# 🎨 Gradio Interface
 demo = gr.Interface(
-    fn=recommend_gifts,
-    inputs="text",
+    fn=recommend_gifts, 
+    inputs="text", 
     outputs="json",
     title="🎁 AI Gift Recommender",
     description="Enter details about the person you are buying a gift for, and get personalized suggestions with shopping links!",