Spaces:
Sleeping
Sleeping
from transformers import pipeline | |
from transformers import AutoTokenizer, AutoModelForSequenceClassification | |
import re | |
import numpy as np | |
from typing import Dict, List, Tuple | |
class GiftRecommender: | |
def __init__(self): | |
# Initialize NER pipeline for entity extraction | |
self.ner = pipeline("ner", model="dbmdz/bert-large-cased-finetuned-conll03-english") | |
# Initialize sentiment analysis pipeline | |
self.sentiment = pipeline("sentiment-analysis") | |
# Initialize zero-shot classification for interest categorization | |
self.zero_shot = pipeline("zero-shot-classification") | |
# Define interest categories | |
self.interest_categories = [ | |
"art", "music", "sports", "technology", "reading", | |
"travel", "cooking", "gaming", "fashion", "outdoor activities" | |
] | |
# Initialize gift rules database | |
self.gift_rules = { | |
"art": ["art supplies set", "digital drawing tablet", "museum membership", "art classes"], | |
"music": ["wireless headphones", "concert tickets", "vinyl records", "music streaming subscription"], | |
"sports": ["fitness tracker", "sports equipment", "team merchandise", "gym membership"], | |
"technology": ["smart gadgets", "latest electronics", "tech accessories", "coding courses"], | |
"reading": ["e-reader", "book subscription", "rare book editions", "bookstore gift card"], | |
"travel": ["travel accessories", "language learning subscription", "travel guides", "luggage"], | |
"cooking": ["cooking classes", "premium ingredients set", "kitchen gadgets", "cookbook collection"], | |
"gaming": ["gaming console", "gaming accessories", "game subscription", "collectible items"], | |
"fashion": ["designer accessories", "fashion subscription box", "custom jewelry", "premium clothing"], | |
"outdoor": ["camping gear", "hiking equipment", "outdoor experiences", "adventure gear"] | |
} | |
def extract_age(self, text: str) -> int: | |
"""Extract age from text using regex.""" | |
age_pattern = r'\b(\d{1,2})\s*-?\s*years?\s*-?\s*old\b|\b(\d{1,2})\b' | |
matches = re.findall(age_pattern, text) | |
if matches: | |
# Return the first number found | |
age = next(int(num) for nums in matches for num in nums if num) | |
return age if 0 < age < 120 else None | |
return None | |
def extract_gender(self, text: str) -> str: | |
"""Extract gender from text using keywords.""" | |
text = text.lower() | |
gender_indicators = { | |
'female': ['she', 'her', 'sister', 'girlfriend', 'wife', 'daughter', 'mom', 'mother'], | |
'male': ['he', 'him', 'brother', 'boyfriend', 'husband', 'son', 'dad', 'father'] | |
} | |
for gender, indicators in gender_indicators.items(): | |
if any(indicator in text for indicator in indicators): | |
return gender | |
return "unknown" | |
def extract_interests(self, text: str) -> List[Dict]: | |
"""Extract and categorize interests using zero-shot classification.""" | |
# First, extract potential interest phrases | |
interest_pattern = r'loves?\s+([^,.]+)|\blikes?\s+([^,.]+)' | |
matches = re.findall(interest_pattern, text.lower()) | |
interests = [] | |
for match in matches: | |
phrase = next(m for m in match if m) | |
# Classify the interest into predefined categories | |
result = self.zero_shot( | |
phrase, | |
candidate_labels=self.interest_categories, | |
multi_label=False | |
) | |
# Get sentiment score for the interest | |
sentiment_score = self.sentiment(phrase)[0] | |
interests.append({ | |
'phrase': phrase, | |
'category': result['labels'][0], | |
'confidence': result['scores'][0], | |
'sentiment': sentiment_score['label'], | |
'sentiment_score': sentiment_score['score'] | |
}) | |
return interests | |
def extract_dislikes(self, text: str) -> List[str]: | |
"""Extract dislikes from text.""" | |
dislike_pattern = r'hates?\s+([^,.]+)|dislikes?\s+([^,.]+)' | |
matches = re.findall(dislike_pattern, text.lower()) | |
return [next(m for m in match if m) for match in matches] | |
def get_gift_recommendations(self, text: str) -> Dict: | |
"""Process text and return gift recommendations.""" | |
# Extract basic information | |
age = self.extract_age(text) | |
gender = self.extract_gender(text) | |
interests = self.extract_interests(text) | |
dislikes = self.extract_dislikes(text) | |
# Generate recommendations based on interests | |
recommendations = [] | |
for interest in interests: | |
category = interest['category'] | |
if category in self.gift_rules: | |
# Weight recommendations by sentiment and confidence | |
weight = interest['confidence'] * (2 if interest['sentiment'] == 'POSITIVE' else 1) | |
recommendations.extend([ | |
{ | |
'gift': gift, | |
'category': category, | |
'weight': weight, | |
'reason': f"Based on their interest in {interest['phrase']}" | |
} | |
for gift in self.gift_rules[category] | |
]) | |
# Sort recommendations by weight | |
recommendations.sort(key=lambda x: x['weight'], reverse=True) | |
return { | |
'profile': { | |
'age': age, | |
'gender': gender, | |
'interests': interests, | |
'dislikes': dislikes | |
}, | |
'recommendations': recommendations[:5] # Return top 5 recommendations | |
} | |
def format_recommendations(self, results: Dict) -> str: | |
"""Format the recommendations into a readable string.""" | |
output = [] | |
output.append("π Gift Recommendations\n") | |
profile = results['profile'] | |
output.append(f"Profile Summary:") | |
output.append(f"- Age: {profile['age'] or 'Unknown'}") | |
output.append(f"- Gender: {profile['gender'].title()}") | |
output.append("- Interests: " + ", ".join(i['phrase'] for i in profile['interests'])) | |
if profile['dislikes']: | |
output.append("- Dislikes: " + ", ".join(profile['dislikes'])) | |
output.append("\nTop Recommendations:") | |
for i, rec in enumerate(results['recommendations'], 1): | |
output.append(f"{i}. {rec['gift']}") | |
output.append(f" β’ {rec['reason']}") | |
return "\n".join(output) | |
# Example usage | |
if __name__ == "__main__": | |
recommender = GiftRecommender() | |
# Example input | |
text = """I'm looking for a gift for my 25-year-old sister. | |
She loves painting and traveling, especially in Japan. | |
She hates loud noises and doesn't like spicy food.""" | |
results = recommender.get_gift_recommendations(text) | |
formatted_output = recommender.format_recommendations(results) | |
print(formatted_output) | |