Update app.py

app.py

@@ -1,125 +1,75 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-import torch
+from transformers import pipeline
 import requests
-from typing import List, Dict
 
-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(
+def extract_interests(text):
+    """Extract interests using a smaller, more efficient model"""
+    try:
+        # Using GPT-2 small model which is more lightweight
+        classifier = pipeline(
             "text-generation",
-            model=self.model,
-            tokenizer=self.tokenizer,
+            model="gpt2",
+            max_length=50,
             device_map="auto"
         )
-    
-    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)
         
+        prompt = f"List 3 interests from: {text}\nInterests:"
+        response = classifier(prompt, max_new_tokens=30, num_return_sequences=1)
+        
+        # Clean up the response
+        interests = response[0]['generated_text'].split('Interests:')[-1].strip()
+        return [i.strip() for i in interests.split(',') if i.strip()][:3]
+    except Exception as e:
+        print(f"Error in extraction: {str(e)}")
+        return ["general gifts"]
+
+def search_gifts(interests):
+    """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}"
+    }
+
+def recommend_gifts(text):
+    """Main recommendation function"""
+    if not text:
         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}"
+            "error": "Please enter a description."
         }
     
-    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 = extract_interests(text)
         
-        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)}"
-            }
+        # Get shopping links
+        links = search_gifts(interests)
+        
+        return {
+            "Identified Interests": interests,
+            "Shopping Links": links
+        }
+    except Exception as e:
+        return {
+            "error": f"An error occurred: {str(e)}"
+        }
 
-# 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
+# Create Gradio interface
+demo = gr.Interface(
+    fn=recommend_gifts,
+    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!",
+    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"],
+    ]
+)
 
 if __name__ == "__main__":
-    demo = create_gradio_interface()
-    demo.launch(share=True)  # Enable sharing for testing
+    demo.launch()