forked from zephyrtronium/robot
/
main.go
113 lines (103 loc) · 2.75 KB
/
main.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
101
102
103
104
105
106
107
108
109
110
111
112
113
/*
Copyright (C) 2020 Branden J Brown
This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.
You should have received a copy of the GNU General Public License
along with this program. If not, see <https://www.gnu.org/licenses/>.
*/
// robot-convert imports old Robot serialization files to the new brain format.
package main
import (
"context"
"database/sql"
"encoding/json"
"flag"
"io/ioutil"
"log"
"strings"
"github.com/zephyrtronium/robot/brain"
)
func main() {
var me, tag, in, out string
flag.StringVar(&me, "me", "", "bot username")
flag.StringVar(&tag, "tag", "", "input tag")
flag.StringVar(&in, "in", "", "input json database")
flag.StringVar(&out, "out", "", "output SQLite3 database")
flag.Parse()
if me == "" {
log.Fatalln("-me is required")
}
if tag == "" {
log.Fatalln("-tag is required")
}
if in == "" {
log.Fatalln("-in is required")
}
if out == "" {
log.Fatalln("-out is required")
}
b, err := ioutil.ReadFile(in)
if err != nil {
log.Fatalln(err)
}
var chain map[string][]string
if err := json.Unmarshal(b, &chain); err != nil {
log.Fatalln(err)
}
var order int
for p := range chain {
order = len(strings.Fields(p)) // easiest consistent measure
break
}
ctx := context.Background()
br, err := brain.Configure(ctx, out, me, order)
if err != nil {
log.Fatalln(err)
}
defer br.Close()
// SQLite: Always use WAL journaling. It can be orders of magnitude faster.
if _, err := br.Exec(ctx, `PRAGMA journal_mode=WAL`); err != nil {
log.Println("unable to set journal mode to WAL:", err)
}
pre := make([]sql.NullString, order)
var n int64
for p, s := range chain {
f := strings.Fields(p)
for i, v := range f {
if v == "\x01" {
pre[i] = sql.NullString{}
} else {
pre[i] = sql.NullString{String: v, Valid: true}
}
}
for _, w := range s {
switch w {
case "\x01":
continue // ??
case "\x00":
r, err := br.LearnTuple(ctx, tag, pre, sql.NullString{})
if err != nil {
log.Printf("error learning %v -> NULL: %v", pre, err)
continue
}
k, _ := r.RowsAffected()
n += k
default:
r, err := br.LearnTuple(ctx, tag, pre, sql.NullString{String: w, Valid: true})
if err != nil {
log.Printf("error learning %v -> %s: %v", pre, w, err)
continue
}
k, _ := r.RowsAffected()
n += k
}
}
}
log.Println("learned", n, "chains")
}