-
Notifications
You must be signed in to change notification settings - Fork 23
/
file-util.go
246 lines (211 loc) · 7.01 KB
/
file-util.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
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
package tokenizer
import (
"fmt"
"io"
"log"
"net/http"
"os"
"path"
"strconv"
"strings"
)
// This file provides functions to work with local dataset cache, ...
// Copied from `transformer` package
const (
WeightName = "pytorch_model.gt"
ConfigName = "config.json"
TokenizerName = "tokenizer.json"
// NOTE. URL form := `$HFpath/ModelName/resolve/main/WeightName`
HFpath = "https://huggingface.co"
)
var (
DUMMY_INPUT [][]int64 = [][]int64{
{7, 6, 0, 0, 1},
{1, 2, 3, 0, 0},
{0, 0, 0, 4, 5},
}
)
// CachedPath resolves and caches data based on input string, then returns fullpath to the cached data.
//
// Parameters:
// - `modelNameOrPath`: model name e.g., "bert-base-uncase" or path to directory contains model/config files.
// - `fileName`: model or config file name. E.g., "pytorch_model.py", "config.json"
//
// CachedPath does several things consequently:
// 1. Resolves input string to a fullpath cached filename candidate.
// 2. Check it at `CachedPath`, if exists, then return the candidate. If not
// 3. Retrieves and Caches data to `CachedPath` and returns path to cached data
//
// NOTE. default `CachedDir` is at "{$HOME}/.cache/transformer"
// Custom `CachedDir` can be changed by setting with environment `GO_TRANSFORMER`
func CachedPath(modelNameOrPath, fileName string) (resolvedPath string, err error) {
// Resolves to "candidate" filename at `CacheDir`
cachedFileCandidate := fmt.Sprintf("%s/%s/%s", CachedDir, modelNameOrPath, fileName)
// 1. Cached candidate file exists
if _, err := os.Stat(cachedFileCandidate); err == nil {
return cachedFileCandidate, nil
}
// 2. If valid fullpath to local file, caches it and return cached filename
filepath := fmt.Sprintf("%s/%s", modelNameOrPath, fileName)
if _, err := os.Stat(filepath); err == nil {
err := copyFile(filepath, cachedFileCandidate)
if err != nil {
err := fmt.Errorf("CachedPath() failed at copying file: %w", err)
return "", err
}
return cachedFileCandidate, nil
}
// 3. Cached candidate file NOT exist. Try to download it and save to `CachedDir`
url := fmt.Sprintf("%s/%s/resolve/main/%s", HFpath, modelNameOrPath, fileName)
// url := fmt.Sprintf("%s/%s/raw/main/%s", HFpath, modelNameOrPath, fileName)
if isValidURL(url) {
if _, err := http.Get(url); err == nil {
err := downloadFile(url, cachedFileCandidate)
if err != nil {
err = fmt.Errorf("CachedPath() failed at trying to download file: %w", err)
return "", err
}
return cachedFileCandidate, nil
} else {
err = fmt.Errorf("CachedPath() failed: Unable to parse '%v' as a URL or as a local path.\n", url)
return "", err
}
}
// Not resolves
err = fmt.Errorf("CachedPath() failed: Unable to parse '%v' as a URL or as a local path.\n", url)
return "", err
}
func isValidURL(url string) bool {
// TODO: implement
return true
}
// downloadFile downloads file from URL and stores it in local filepath.
// It writes to the destination file as it downloads it, without loading
// the entire file into memory. An `io.TeeReader` is passed into Copy()
// to report progress on the download.
func downloadFile(url string, filepath string) error {
// Create path if not existing
dir := path.Dir(filepath)
filename := path.Base(filepath)
if _, err := os.Stat(dir); os.IsNotExist(err) {
if err := os.MkdirAll(dir, 0755); err != nil {
log.Fatal(err)
}
}
// Create the file with .tmp extension, so that we won't overwrite a
// file until it's downloaded fully
out, err := os.Create(filepath + ".tmp")
if err != nil {
return err
}
defer out.Close()
// Get the data
resp, err := http.Get(url)
if err != nil {
return err
}
defer resp.Body.Close()
// Check server response
if resp.StatusCode != http.StatusOK {
err := fmt.Errorf("bad status: %s(%v)", resp.Status, resp.StatusCode)
if resp.StatusCode == 404 {
// if filename == "rust_model.ot" {
// msg := fmt.Sprintf("model weight file not found. That means a compatible pretrained model weight file for Go is not available.\n")
// msg = msg + fmt.Sprintf("You might need to manually convert a 'pytorch_model.bin' for Go. ")
// msg = msg + fmt.Sprintf("See tutorial at: 'example/convert'")
// err = fmt.Errorf(msg)
// } else {
// err = fmt.Errorf("download file not found: %q for downloading", url)
// }
err = fmt.Errorf("download file not found: %q for downloading", url)
} else {
err = fmt.Errorf("download file failed: %q", url)
}
return err
}
// the total file size to download
size, _ := strconv.Atoi(resp.Header.Get("Content-Length"))
downloadSize := uint64(size)
// Create our bytes counter and pass it to be used alongside our writer
counter := &writeCounter{FileSize: downloadSize}
_, err = io.Copy(out, io.TeeReader(resp.Body, counter))
if err != nil {
return err
}
fmt.Printf("\r%s... %s/%s completed", filename, byteCountIEC(counter.Total), byteCountIEC(counter.FileSize))
// The progress use the same line so print a new line once it's finished downloading
fmt.Println()
// Rename the tmp file back to the original file
err = os.Rename(filepath+".tmp", filepath)
if err != nil {
return err
}
return nil
}
// writeCounter counts the number of bytes written to it. By implementing the Write method,
// it is of the io.Writer interface and we can pass this into io.TeeReader()
// Every write to this writer, will print the progress of the file write.
type writeCounter struct {
Total uint64
FileSize uint64
}
func (wc *writeCounter) Write(p []byte) (int, error) {
n := len(p)
wc.Total += uint64(n)
wc.printProgress()
return n, nil
}
// PrintProgress prints the progress of a file write
func (wc writeCounter) printProgress() {
// Clear the line by using a character return to go back to the start and remove
// the remaining characters by filling it with spaces
fmt.Printf("\r%s", strings.Repeat(" ", 50))
// Return again and print current status of download
fmt.Printf("\rDownloading... %s/%s", byteCountIEC(wc.Total), byteCountIEC(wc.FileSize))
}
// byteCountIEC converts bytes to human-readable string in binary (IEC) format.
func byteCountIEC(b uint64) string {
const unit = 1024
if b < unit {
return fmt.Sprintf("%d B", b)
}
div, exp := uint64(unit), 0
for n := b / unit; n >= unit; n /= unit {
div *= unit
exp++
}
return fmt.Sprintf("%.1f %ciB",
float64(b)/float64(div), "KMGTPE"[exp])
}
func copyFile(src, dst string) error {
sourceFileStat, err := os.Stat(src)
if err != nil {
return err
}
if !sourceFileStat.Mode().IsRegular() {
return fmt.Errorf("%s is not a regular file", src)
}
source, err := os.Open(src)
if err != nil {
return err
}
defer source.Close()
destination, err := os.Create(dst)
if err != nil {
return err
}
defer destination.Close()
_, err = io.Copy(destination, source)
return err
}
// CleanCache removes all files cached in transformer cache directory `CachedDir`.
//
// NOTE. custom `CachedDir` can be changed by setting environment `GO_TRANSFORMER`
func CleanCache() error {
err := os.RemoveAll(CachedDir)
if err != nil {
err = fmt.Errorf("CleanCache() failed: %w", err)
return err
}
return nil
}