-
Notifications
You must be signed in to change notification settings - Fork 0
/
model.go
100 lines (86 loc) · 2.23 KB
/
model.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
package server
import (
"bufio"
"fmt"
"os"
"strings"
"github.com/jbochi/facts/vectormodel"
"github.com/kshedden/gonpy"
)
type (
// Model is the struct that handles recommendations
Model struct {
vm *vectormodel.VectorModel
repositories []string
repositoryIDs map[string]int
}
// RepositoryScore is a pair of repo / score
RepositoryScore struct {
Repository string
Score float64
}
)
// ReadModel returns a VectorModel from given file path
func ReadModel(path string) (*Model, error) {
confidence := 3.0
regularization := 0.001
rdr, err := gonpy.NewFileReader(path + "item_factors.npy")
if err != nil {
return nil, fmt.Errorf("Unable to read data: %v", err)
}
nRepositories, nFactors := rdr.Shape[0], rdr.Shape[1]
data, err := rdr.GetFloat64()
if err != nil {
return nil, fmt.Errorf("Unable to parse data: %v", err)
}
docs := make(map[int][]float64)
for i := 0; i < nRepositories; i++ {
docs[i] = data[i*nFactors : (i+1)*nFactors]
}
vm, err := vectormodel.NewVectorModel(docs, confidence, regularization)
if err != nil {
return nil, err
}
f, err := os.Open(path + "items.csv")
if err != nil {
return nil, fmt.Errorf("Unable to open items.csv: %v", err)
}
repositories := make([]string, nRepositories)
repositoryIDs := map[string]int{}
reader := bufio.NewReader(f)
for i := 0; i < rdr.Shape[0]; i++ {
line, err := reader.ReadString('\n')
if err != nil {
return nil, fmt.Errorf("Unable to read line of file: %v", err)
}
repo := strings.TrimRight(line, "\n")
repositories[i] = repo
repositoryIDs[repo] = i
}
m := &Model{
vm: vm,
repositories: repositories,
repositoryIDs: repositoryIDs,
}
return m, nil
}
// Recommend returns a list of recommended repositories
func (m *Model) Recommend(items []string, n int) ([]RepositoryScore, error) {
seenDocs := map[int]bool{}
for _, repo := range items {
repoID, ok := m.repositoryIDs[repo]
if ok {
seenDocs[repoID] = true
}
}
scores, err := m.vm.Recommend(&seenDocs, n)
if err != nil {
return nil, err
}
results := []RepositoryScore{}
for _, score := range scores {
result := RepositoryScore{m.repositories[score.DocumentID], score.Score}
results = append(results, result)
}
return results, nil
}