Gift-Recommender / product_recommender.py
noddysnots's picture
Update product_recommender.py
dc7e417 verified
raw
history blame
1.51 kB
from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoModelForTokenClassification
from sentence_transformers import SentenceTransformer
import torch
import numpy as np
from typing import Dict, List, Optional
import json
class ProductRecommender:
def __init__(self):
self.model_analyzer = MultiModelAnalyzer()
def get_recommendations(self, query: str, product_database: List[Dict]) -> List[Dict]:
try:
query_analysis = self.model_analyzer.analyze_text(query)
# For testing, return a simple recommendation
return [{
"name": "Test Product",
"price": "β‚Ή999",
"category": "test",
"similarity": 0.95
}]
except Exception as e:
print(f"Error in recommendations: {str(e)}")
return []
class MultiModelAnalyzer:
def __init__(self):
try:
self.category_model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased")
self.category_tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
self.semantic_model = SentenceTransformer('paraphrase-MiniLM-L6-v2')
except Exception as e:
print(f"Error initializing models: {str(e)}")
def analyze_text(self, text: str) -> Dict:
return {
"category": "test",
"embedding": np.zeros(10),
"features": ["test"]
}