FranKGraphBench 0.0.1a0__py2.py3-none-any.whl

framework/recommender/models/distMult/model.py

@@ -0,0 +1,187 @@
+from framework.dataloader.graph.node import UserNode, ItemNode
+from ...recommender import Recommender
+
+from pykeen.pipeline import pipeline
+from pykeen.triples import TriplesFactory
+from pykeen.datasets import Dataset
+from sklearn.neighbors import NearestNeighbors
+
+
+import numpy as np
+import torch
+import typing as t
+
+"""
+    Recommender base class
+"""
+
+
+class DistMult(Recommender):
+    def __init__(
+        self,
+        config: dict,
+        embedding_dim: int = 50,
+        entity_initializer=None,
+        entity_constrainer=None,
+        relation_initializer=None,
+        regularizer=None,
+        regularizer_kwargs=None,
+        epochs: int = 5,
+        seed: int = 42,
+        evaluation_fallback: bool = True,
+        all_recs: bool = False,
+        triples: str = "all",
+    ):
+        super().__init__(config)
+        self.embedding_dim = embedding_dim
+        self.entity_initializer = entity_initializer
+        self.entity_constrainer = entity_constrainer
+        self.relation_initializer = relation_initializer
+        self.regularizer = regularizer
+        self.regularizer_kwargs = regularizer_kwargs
+        self.epochs = epochs
+        self.seed = seed
+        self.evaluation_fallback = evaluation_fallback
+        self.all_recs = all_recs
+        self.triples = triples
+        self._triples = None
+        self._embedding = {}
+
+    def name(self):
+        text = "DistMult based model + cosine similarity"
+        text += f";embedding_dim={self.embedding_dim};entity_initializer={self.entity_initializer};entity_constrainer={self.entity_constrainer};relation_initializer={self.relation_initializer};regularizer={self.regularizer};epochs={self.epochs};seed={self.seed};triples={self.triples}"
+        return text
+
+    def train(self, G_train, ratings_train):
+        self.G_train = G_train
+        if self.triples == "all":
+            self._triples = G_train.get_all_triples()
+        elif self.triples == "ratings":
+            self._triples = G_train.get_ratings_triples()
+        self.fit()
+
+    def get_recommendations(self, k: int = 5) -> t.Dict[UserNode, t.List[ItemNode]]:
+        """
+        :param k: cutoff recommendation (int)
+        :return dict of recommendations for each user
+            of type {user1: [item1, item2]}
+        """
+        users = list(self.G_train.get_user_nodes())
+        items = list(self.G_train.get_item_nodes())
+
+        users_embeddings = np.array([self._embedding[user] for user in users])
+        items_embeddings = np.array([self._embedding[item] for item in items])
+
+        n_neighbors = self._get_n_neighbors(users, items, k)
+        knn = NearestNeighbors(n_neighbors=n_neighbors, metric="cosine")
+        knn.fit(items_embeddings)
+        rec_indices = knn.kneighbors(users_embeddings, return_distance=False)
+
+        recommendations = {}
+        for user_idx, user in enumerate(users):
+            rated_items = self.G_train.get_user_rated_items(user)
+            recs = []
+            for item_idx in rec_indices[user_idx]:
+                item = items[item_idx]
+                if item not in rated_items:
+                    recs.append(item)
+                    if not self.all_recs and len(recs) == k:
+                        break
+
+            recommendations[user] = recs
+
+        return recommendations
+
+    def get_user_recommendation(self, user: UserNode, k: int = 5):
+        raise NotImplementedError(
+            "Override get_user_recommendation() for your model subclass."
+        )
+
+    def fit(self):
+        triples = TriplesFactory.from_labeled_triples(
+            self._triples[["head", "relation", "tail"]].values,
+            create_inverse_triples=False,
+            entity_to_id=None,
+            relation_to_id=None,
+            compact_id=False,
+            filter_out_candidate_inverse_relations=True,
+            metadata=None,
+        )
+
+        dataset = Dataset.from_tf(triples, ratios=[0.95, 0.05, 0.0])
+        entity_to_id = dataset.entity_to_id
+
+        result = pipeline(
+            dataset=dataset,
+            model="DistMult",
+            model_kwargs=dict(
+                embedding_dim=self.embedding_dim,
+                entity_initializer=self.entity_initializer,
+                entity_constrainer=self.entity_constrainer,
+                relation_initializer=self.relation_initializer,
+                regularizer=self.regularizer,
+                regularizer_kwargs=self.regularizer_kwargs,
+            ),
+            epochs = self.epochs,
+            random_seed = self.seed,
+            evaluation_fallback = self.evaluation_fallback,
+        )
+        model = result.model
+
+        users_iter = self.G_train.get_user_nodes()
+        items_iter = self.G_train.get_item_nodes()
+
+        users = list(users_iter)
+        items = list(items_iter)
+
+        # deal with entities that might have not been embedded
+        users_indices, users_no_embedding = [], []
+        for user in users_iter:
+            try:
+                users_indices.append(entity_to_id[user.__str__()])
+            except:
+                users_no_embedding.append(user)
+                users.remove(user)
+                continue
+        users_indices = torch.LongTensor(users_indices)
+
+        items_indices, items_no_embedding = [], []
+        for item in items_iter:
+            try:
+                items_indices.append(entity_to_id[item.__str__()])
+            except:
+                items_no_embedding.append(item)
+                items.remove(item)
+                continue
+        items_indices = torch.LongTensor(items_indices)
+
+        users_embeddings = (model.entity_representations[0](indices=users_indices).detach().cpu().numpy())
+        items_embeddings = (model.entity_representations[0](indices=items_indices).detach().cpu().numpy())
+
+        if len(users_no_embedding) > 0:
+            users_embeddings = np.concatenate((users_embeddings, np.zeros((len(users_no_embedding), self.embedding_dim), dtype=users_embeddings.dtype)), axis=0)
+            users = users + users_no_embedding
+
+        if len(items_no_embedding) > 0:
+            items_embeddings = np.concatenate((items_embeddings, np.zeros((len(items_no_embedding), self.embedding_dim), dtype=items_embeddings.dtype)), axis=0)
+            items = items + items_no_embedding
+
+        for index, user in enumerate(users):
+            self._embedding[user] = users_embeddings[index]
+
+        for index, item in enumerate(items):
+            self._embedding[item] = items_embeddings[index]
+
+    def _get_n_neighbors(self, users, items, top_k):
+        n_neighbors = 0
+        if self.all_recs:
+            n_neighbors = len(items)
+        else:
+            max_recs = 0
+            for user in users:
+                rated_items = self.G_train.get_user_rated_items(user)
+                if len(rated_items) > max_recs:
+                    max_recs = len(rated_items)
+            n_neighbors = min(max_recs + top_k, len(items))
+
+        return n_neighbors

framework/recommender/models/ePHEN/model.py

@@ -0,0 +1,159 @@
+from framework.dataloader.graph.node import UserNode, ItemNode
+from ...recommender import Recommender
+
+from ...model2class import model2class
+
+from sklearn.neighbors import NearestNeighbors
+from sentence_transformers import SentenceTransformer
+
+import numpy as np
+import pandas as pd
+import random
+import importlib
+
+from tqdm import tqdm
+from copy import deepcopy
+
+import os
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+"""
+    Recommender base class
+"""
+
+class EPHEN(Recommender):
+    def __init__(
+        self,
+        config: dict,
+        embedding_model: str = "sentence-transformers/distiluse-base-multilingual-cased-v2",
+        embedding_model_kwargs: dict = None,
+        embed_with: str = "movie_title", # 'graph' for graph embedding model or column for text embedding model
+        iterations: int = 30,
+        mi: float = 0.85,
+        all_recs: bool = False,
+    ):
+        super().__init__(config)
+        self.embedding_model = embedding_model
+        self.embedding_model_kwargs = embedding_model_kwargs
+        self.embed_with = embed_with
+        self.iterations = iterations
+        self.mi = mi
+        self.all_recs = all_recs
+        self._embedding = {}
+
+    def name(self):
+        text = "EPHEN based model + cosine similarity"
+        text += f";embedding_model={self.embedding_model};embedding_model_kwargs={self.embedding_model_kwargs};embed_with={self.embed_with};iterations={self.iterations};mi={self.mi}."
+        return text
+
+    def train(self, G_train, ratings_train):
+        self.G_train = G_train
+        self.ratings_train = ratings_train
+        self.fit()
+
+    def get_recommendations(self, k: int = 5):
+        # Set doesnt guarantee the order of set elements
+        users = list(self.G_train.get_user_nodes())
+        items = list(self.G_train.get_item_nodes())
+
+        users_embeddings = np.array([self._embedding[user] for user in users])
+        items_embeddings = np.array([self._embedding[item] for item in items])
+
+        n_neighbors = self._get_n_neighbors(users, items, k)
+        knn = NearestNeighbors(n_neighbors=n_neighbors, metric='cosine')
+        knn.fit(items_embeddings)
+        rec_indices = knn.kneighbors(users_embeddings, return_distance=False)
+
+        recommendations = {}
+        for user_idx, user in enumerate(users):
+            rated_items = self.G_train.get_user_rated_items(user)
+            recs = []
+            for item_idx in rec_indices[user_idx]:
+                item = items[item_idx]
+                if item not in rated_items:
+                    recs.append(item)
+                    if not self.all_recs and len(recs) == k:
+                        break
+            
+            recommendations[user] = recs
+        
+        return recommendations
+
+    def get_user_recommendation(self, user: UserNode, k: int = 5):
+        raise NotImplementedError(
+            "Override get_user_recommendation() for your model subclass."
+        )
+
+    def fit(self):
+        if self.embed_with == "graph":
+            self._generate_graph_init_embeddings()
+        else:
+            self._generate_text_init_embeddings()
+        self._regularization()
+
+    def _get_n_neighbors(self, users, items, top_k):
+        n_neighbors = 0
+        if self.all_recs:
+            n_neighbors = len(items)
+        else:
+            max_recs = 0
+            for user in users:
+                rated_items = self.G_train.get_user_rated_items(user)
+                if len(rated_items) > max_recs:
+                    max_recs = len(rated_items)
+            n_neighbors = min(max_recs + top_k, len(items))
+
+        return n_neighbors
+    
+    def _generate_text_init_embeddings(self):
+        model = SentenceTransformer(self.embedding_model)
+        item_properties = pd.DataFrame(self.G_train.get_item_property_edges())
+        filtered_properties = item_properties[1][item_properties[1].apply(lambda x: x.get_property_type() == self.embed_with)].reset_index(drop=True)
+        self.init_embeddings = {}
+        for property in tqdm(filtered_properties):
+            self.init_embeddings[property] = model.encode(property.id)
+
+    def _generate_graph_init_embeddings(self):
+        module_name = f'framework.recommender.models.{model2class[self.embedding_model]["submodule"]}'
+        class_name = model2class[self.embedding_model]['class']
+
+        model = getattr(importlib.import_module(module_name), class_name)
+        model = model(self.embedding_model_kwargs['config'], **self.embedding_model_kwargs['parameters'])
+
+        model.train(self.G_train, self.ratings_train)
+
+        self.init_embeddings = model._embedding
+
+    def _regularization(self):
+        dim = len(self.init_embeddings[list(self.init_embeddings.keys())[0]])
+        self._embedding = deepcopy(self.init_embeddings)
+        for node in self.G_train.nodes():
+            if node not in self._embedding:
+                self._embedding[node] = np.array([0.0] * dim)
+        nodes = list(self.G_train.nodes())
+        pbar = tqdm(range(0, self.iterations))
+        for iteration in pbar:
+            random.shuffle(nodes)
+            energy = 0.0
+            for node in nodes:
+                f_new = np.array([0.0] * dim)
+                f_old = np.array(self._embedding[node])
+                sum_w = 0.0
+                for neighbor in self.G_train.neighbors(node):
+                    w = 1.0
+                    if 'weight' in self.G_train[node][neighbor]:
+                        w = self.G_train[node][neighbor]['weight']
+                    w /= np.sqrt(self.G_train.degree[neighbor])
+                    f_new = f_new + w * f_old
+                    sum_w = sum_w + w
+                if sum_w == 0.0: 
+                    sum_w = 1.0
+                f_new /= sum_w
+                self._embedding[node] = f_new
+                if node in self.init_embeddings:
+                    self._embedding[node] = self.init_embeddings[node] * \
+                        self.mi + self._embedding[node] * (1.0 - self.mi)
+                energy = energy + np.linalg.norm(f_new-f_old)
+            iteration = iteration + 1
+            message = f"Iteration {iteration} | Energy = {energy}"
+            pbar.set_description(message)

framework/recommender/models/rESCAL/model.py

@@ -0,0 +1,184 @@
+from framework.dataloader.graph.node import UserNode, ItemNode
+from ...recommender import Recommender
+
+from pykeen.pipeline import pipeline
+from pykeen.triples import TriplesFactory
+from pykeen.datasets import Dataset
+from sklearn.neighbors import NearestNeighbors
+
+
+import numpy as np
+import torch
+import typing as t
+
+"""
+    Recommender base class
+"""
+
+
+class RESCAL(Recommender):
+    def __init__(
+        self,
+        config: dict,
+        embedding_dim: int = 50,
+        entity_initializer=None,
+        relation_initializer=None,
+        regularizer=None,
+        regularizer_kwargs=None,
+        epochs: int = 5,
+        seed: int = 42,
+        evaluation_fallback: bool = True,
+        all_recs: bool = False,
+        triples: str = "all",
+    ):
+        super().__init__(config)
+        self.embedding_dim = embedding_dim
+        self.entity_initializer = entity_initializer
+        self.relation_initializer = relation_initializer
+        self.regularizer = regularizer
+        self.regularizer_kwargs = regularizer_kwargs
+        self.epochs = epochs
+        self.seed = seed
+        self.evaluation_fallback = evaluation_fallback
+        self.all_recs = all_recs
+        self.triples = triples
+        self._triples = None
+        self._embedding = {}
+
+    def name(self):
+        text = "RESCAL based model + cosine similarity"
+        text += f";embedding_dim={self.embedding_dim};entity_initializer={self.entity_initializer};relation_initializer={self.relation_initializer};regularizer={self.regularizer};epochs={self.epochs};seed={self.seed};triples={self.triples}"
+        return text
+
+    def train(self, G_train, ratings_train):
+        self.G_train = G_train
+        if self.triples == "all":
+            self._triples = G_train.get_all_triples()
+        elif self.triples == "ratings":
+            self._triples = G_train.get_ratings_triples()
+        self.fit()
+
+    def get_recommendations(self, k: int = 5) -> t.Dict[UserNode, t.List[ItemNode]]:
+        """
+        :param k: cutoff recommendation (int)
+        :return dict of recommendations for each user
+            of type {user1: [item1, item2]}
+        """
+        users = list(self.G_train.get_user_nodes())
+        items = list(self.G_train.get_item_nodes())
+
+        users_embeddings = np.array([self._embedding[user] for user in users])
+        items_embeddings = np.array([self._embedding[item] for item in items])
+
+        n_neighbors = self._get_n_neighbors(users, items, k)
+        knn = NearestNeighbors(n_neighbors=n_neighbors, metric="cosine")
+        knn.fit(items_embeddings)
+        rec_indices = knn.kneighbors(users_embeddings, return_distance=False)
+
+        recommendations = {}
+        for user_idx, user in enumerate(users):
+            rated_items = self.G_train.get_user_rated_items(user)
+            recs = []
+            for item_idx in rec_indices[user_idx]:
+                item = items[item_idx]
+                if item not in rated_items:
+                    recs.append(item)
+                    if not self.all_recs and len(recs) == k:
+                        break
+
+            recommendations[user] = recs
+
+        return recommendations
+
+    def get_user_recommendation(self, user: UserNode, k: int = 5):
+        raise NotImplementedError(
+            "Override get_user_recommendation() for your model subclass."
+        )
+
+    def fit(self):
+        triples = TriplesFactory.from_labeled_triples(
+            self._triples[["head", "relation", "tail"]].values,
+            create_inverse_triples=False,
+            entity_to_id=None,
+            relation_to_id=None,
+            compact_id=False,
+            filter_out_candidate_inverse_relations=True,
+            metadata=None,
+        )
+
+        dataset = Dataset.from_tf(triples, ratios=[0.95, 0.05, 0.0])
+        entity_to_id = dataset.entity_to_id
+
+        result = pipeline(
+            dataset=dataset,
+            model="RESCAL",
+            model_kwargs=dict(
+                embedding_dim=self.embedding_dim,
+                entity_initializer=self.entity_initializer,
+                relation_initializer=self.relation_initializer,
+                regularizer=self.regularizer,
+                regularizer_kwargs=self.regularizer_kwargs,
+            ),
+            epochs = self.epochs,
+            random_seed = self.seed,
+            evaluation_fallback = self.evaluation_fallback,
+        )
+        model = result.model
+
+        users_iter = self.G_train.get_user_nodes()
+        items_iter = self.G_train.get_item_nodes()
+
+        users = list(users_iter)
+        items = list(items_iter)
+
+        # deal with entities that might have not been embedded
+        users_indices, users_no_embedding = [], []
+        for user in users_iter:
+            try:
+                users_indices.append(entity_to_id[user.__str__()])
+            except:
+                users_no_embedding.append(user)
+                users.remove(user)
+                continue
+        users_indices = torch.LongTensor(users_indices)
+
+        items_indices, items_no_embedding = [], []
+        for item in items_iter:
+            try:
+                items_indices.append(entity_to_id[item.__str__()])
+            except:
+                items_no_embedding.append(item)
+                items.remove(item)
+                continue
+        items_indices = torch.LongTensor(items_indices)
+
+        users_embeddings = (model.entity_representations[0](indices=users_indices).detach().cpu().numpy())
+        items_embeddings = (model.entity_representations[0](indices=items_indices).detach().cpu().numpy())
+
+        if len(users_no_embedding) > 0:
+            users_embeddings = np.concatenate((users_embeddings, np.zeros((len(users_no_embedding), self.embedding_dim), dtype=users_embeddings.dtype)), axis=0)
+            users = users + users_no_embedding
+
+        if len(items_no_embedding) > 0:
+            items_embeddings = np.concatenate((items_embeddings, np.zeros((len(items_no_embedding), self.embedding_dim), dtype=items_embeddings.dtype)), axis=0)
+            items = items + items_no_embedding
+
+        for index, user in enumerate(users):
+            self._embedding[user] = users_embeddings[index]
+
+        for index, item in enumerate(items):
+            self._embedding[item] = items_embeddings[index]
+
+    def _get_n_neighbors(self, users, items, top_k):
+        n_neighbors = 0
+        if self.all_recs:
+            n_neighbors = len(items)
+        else:
+            max_recs = 0
+            for user in users:
+                rated_items = self.G_train.get_user_rated_items(user)
+                if len(rated_items) > max_recs:
+                    max_recs = len(rated_items)
+            n_neighbors = min(max_recs + top_k, len(items))
+
+        return n_neighbors