Merge pull request wzhe06#31 from happy-yuxuan/bugfix

fix bug of generating user emb
Aro710 · Feb 26, 2021 · ec50fa3 · ec50fa3
2 parents 5e27025 + 3eed2f6
commit ec50fa3
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/main/java/com/sparrowrecsys/offline/spark/embedding/Embedding.scala b/src/main/java/com/sparrowrecsys/offline/spark/embedding/Embedding.scala
@@ -62,15 +62,17 @@ object Embedding {
         val userId = user._1
         var userEmb = new Array[Float](embLength)
 
+        var movieCount = 0
         userEmb = user._2.foldRight[Array[Float]](userEmb)((row, newEmb) => {
           val movieId = row.getAs[String]("movieId")
           val movieEmb = word2VecModel.getVectors.get(movieId)
+          movieCount += 1
           if(movieEmb.isDefined){
             newEmb.zip(movieEmb.get).map { case (x, y) => x + y }
           }else{
             newEmb
           }
-        })
+        }).map((x: Float) => x / movieCount)
         userEmbeddings.append((userId,userEmb))
       })