Update app.py

app.py

@@ -1,52 +1,125 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import torch
+import requests
+from typing import List, Dict
 
-# ✅ Use Llama 2 - 7B Chat (Public)
-model_name = "meta-llama/Llama-2-7b-chat-hf"
-
-# 🔹 Authenticate if required (for gated models)
-tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=True)
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
-
-pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=100)
-
-# 🎯 Extract interests from user input
-def extract_interests(text):
-    prompt = f"Extract 3-5 keywords from this: '{text}' (focus on hobbies & products)."
-    response = pipe(prompt)
-    interests = response[0]["generated_text"].replace(prompt, "").strip()
-    return interests.split(", ")
-
-# 🎁 Web search for gift suggestions
-def search_gifts(interests):
-    query = "+".join(interests)
-    return {
-        "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}",
-    }
-
-# 🎯 Main function for gift recommendation
-def recommend_gifts(text):
-    if not text:
-        return "Please enter a description."
+class GiftRecommender:
+    def __init__(self, model_name: str = "deepseek-ai/deepseek-coder-6.7b-base"):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"Using device: {self.device}")
+        
+        # Initialize tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_name,
+            trust_remote_code=True
+        )
+        
+        # Initialize model with appropriate settings
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
+            device_map="auto",
+            trust_remote_code=True,
+            low_cpu_mem_usage=True
+        )
+        
+        # Create pipeline
+        self.generator = pipeline(
+            "text-generation",
+            model=self.model,
+            tokenizer=self.tokenizer,
+            device_map="auto"
+        )
     
-    interests = extract_interests(text)
-    links = search_gifts(interests)
-
-    return {"Predicted Interests": interests, "Gift Suggestions": links}
+    def extract_interests(self, text: str) -> List[str]:
+        """Extract interests from user input."""
+        try:
+            prompt = (
+                f"Extract 3-5 relevant interests or keywords from this text as a comma-separated list. "
+                f"Focus on hobbies and product preferences: '{text}'"
+            )
+            
+            # Generate response with controlled parameters
+            response = self.generator(
+                prompt,
+                max_length=100,
+                num_return_sequences=1,
+                temperature=0.7,
+                top_p=0.95,
+                do_sample=True
+            )
+            
+            # Clean and process the response
+            generated_text = response[0]["generated_text"].replace(prompt, "").strip()
+            interests = [item.strip() for item in generated_text.split(",") if item.strip()]
+            
+            return interests[:5]  # Ensure we return max 5 interests
+            
+        except Exception as e:
+            print(f"Error in interest extraction: {str(e)}")
+            return ["general gifts"]  # Fallback interests
+    
+    def generate_search_links(self, interests: List[str]) -> Dict[str, str]:
+        """Generate shopping links based on interests."""
+        query = "+".join(interests)
+        
+        return {
+            "Amazon India": f"https://www.amazon.in/s?k={query}",
+            "Flipkart": f"https://www.flipkart.com/search?q={query}",
+            "IGP Gifts": f"https://www.igp.com/search?q={query}",
+            "IndiaMart": f"https://dir.indiamart.com/search.mp?ss={query}"
+        }
+    
+    def recommend(self, text: str) -> Dict:
+        """Main recommendation function."""
+        if not text or not text.strip():
+            return {
+                "error": "Please enter a description of the person you're buying a gift for."
+            }
+        
+        try:
+            # Extract interests
+            interests = self.extract_interests(text)
+            
+            # Generate shopping links
+            links = self.generate_search_links(interests)
+            
+            return {
+                "Identified Interests": interests,
+                "Shopping Links": links
+            }
+            
+        except Exception as e:
+            return {
+                "error": f"An error occurred: {str(e)}"
+            }
 
-# 🎨 Gradio UI
-demo = gr.Interface(
-    fn=recommend_gifts, 
-    inputs="text", 
-    outputs="json",
-    title="🎁 AI Gift Recommender (Llama 2 - 7B)",
-    description="Enter details about the person and get personalized gift suggestions with shopping links!",
-)
+# Initialize Gradio interface
+def create_gradio_interface():
+    # Initialize the recommender
+    recommender = GiftRecommender()
+    
+    # Create the interface
+    demo = gr.Interface(
+        fn=recommender.recommend,
+        inputs=gr.Textbox(
+            lines=3,
+            placeholder="Describe the person you're buying a gift for (their interests, hobbies, age, occasion, etc.)"
+        ),
+        outputs=gr.JSON(),
+        title="🎁 AI Gift Recommender",
+        description=(
+            "Enter details about the person you're buying a gift for, and get personalized "
+            "suggestions with shopping links! The AI will analyze their interests and provide relevant recommendations."
+        ),
+        examples=[
+            ["My friend is a 25-year-old software developer who loves gaming and anime"],
+            ["Looking for a gift for my mom who enjoys gardening and cooking"],
+        ]
+    )
+    return demo
 
-# 🚀 Launch Gradio App
 if __name__ == "__main__":
-    demo.launch()
+    demo = create_gradio_interface()
+    demo.launch(share=True)  # Enable sharing for testing