bar.go 3.44 KB
Newer Older
Michael Yang's avatar
Michael Yang committed
1
2
3
4
package progress

import (
	"fmt"
5
	"math"
Michael Yang's avatar
Michael Yang committed
6
7
8
9
10
11
12
13
	"os"
	"strings"
	"time"

	"github.com/jmorganca/ollama/format"
	"golang.org/x/term"
)

14
15
16
17
18
19
type Stats struct {
	rate      int64
	value     int64
	remaining time.Duration
}

Michael Yang's avatar
Michael Yang committed
20
21
22
23
24
25
26
27
28
type Bar struct {
	message      string
	messageWidth int

	maxValue     int64
	initialValue int64
	currentValue int64

	started time.Time
29
30
31

	stats   Stats
	statted time.Time
Michael Yang's avatar
Michael Yang committed
32
33
34
35
36
37
38
39
40
41
42
43
44
}

func NewBar(message string, maxValue, initialValue int64) *Bar {
	return &Bar{
		message:      message,
		messageWidth: -1,
		maxValue:     maxValue,
		initialValue: initialValue,
		currentValue: initialValue,
		started:      time.Now(),
	}
}

45
46
47
48
49
50
51
52
53
54
55
56
57
// formatDuration limits the rendering of a time.Duration to 2 units
func formatDuration(d time.Duration) string {
	if d >= 100*time.Hour {
		return "99h+"
	}

	if d >= time.Hour {
		return fmt.Sprintf("%dh%dm", int(d.Hours()), int(d.Minutes())%60)
	}

	return d.Round(time.Second).String()
}

Michael Yang's avatar
Michael Yang committed
58
59
60
func (b *Bar) String() string {
	termWidth, _, err := term.GetSize(int(os.Stderr.Fd()))
	if err != nil {
61
		termWidth = 80
Michael Yang's avatar
Michael Yang committed
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
	}

	var pre, mid, suf strings.Builder

	if b.message != "" {
		message := strings.TrimSpace(b.message)
		if b.messageWidth > 0 && len(message) > b.messageWidth {
			message = message[:b.messageWidth]
		}

		fmt.Fprintf(&pre, "%s", message)
		if b.messageWidth-pre.Len() >= 0 {
			pre.WriteString(strings.Repeat(" ", b.messageWidth-pre.Len()))
		}

		pre.WriteString(" ")
	}

80
	fmt.Fprintf(&pre, "%3.0f%% ", math.Floor(b.percent()))
81

82
	fmt.Fprintf(&suf, "(%s/%s", format.HumanBytes(b.currentValue), format.HumanBytes(b.maxValue))
Michael Yang's avatar
Michael Yang committed
83

84
	stats := b.Stats()
85
86
87
	rate := stats.rate
	if stats.value > b.initialValue && stats.value < b.maxValue {
		fmt.Fprintf(&suf, ", %s/s", format.HumanBytes(int64(rate)))
88
89
90
91
	}

	fmt.Fprintf(&suf, ")")

92
93
	var timing string
	if stats.value > b.initialValue && stats.value < b.maxValue {
94
		timing = fmt.Sprintf("[%s:%s]", formatDuration(time.Since(b.started)), formatDuration(stats.remaining))
95
	}
Michael Yang's avatar
Michael Yang committed
96

97
	// 44 is the maximum width for the stats on the right of the progress bar
98
99
100
101
	pad := 44 - suf.Len() - len(timing)
	if pad > 0 {
		suf.WriteString(strings.Repeat(" ", pad))
	}
102
	suf.WriteString(timing)
Michael Yang's avatar
Michael Yang committed
103

104
105
	// add 3 extra spaces: 2 boundary characters and 1 space at the end
	f := termWidth - pre.Len() - suf.Len() - 3
Michael Yang's avatar
Michael Yang committed
106
107
	n := int(float64(f) * b.percent() / 100)

108
109
	if f > 0 {
		mid.WriteString("▕")
110
		mid.WriteString(strings.Repeat("█", n))
111
112
113
114
		if f-n > 0 {
			mid.WriteString(strings.Repeat(" ", f-n))
		}
		mid.WriteString("▏")
Michael Yang's avatar
Michael Yang committed
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
	}

	return pre.String() + mid.String() + suf.String()
}

func (b *Bar) Set(value int64) {
	if value >= b.maxValue {
		value = b.maxValue
	}

	b.currentValue = value
}

func (b *Bar) percent() float64 {
	if b.maxValue > 0 {
		return float64(b.currentValue) / float64(b.maxValue) * 100
	}

	return 0
}

136
func (b *Bar) Stats() Stats {
137
	if time.Since(b.statted) < time.Second {
138
		return b.stats
Michael Yang's avatar
Michael Yang committed
139
140
	}

141
142
143
	switch {
	case b.statted.IsZero():
		b.stats = Stats{
Jeffrey Morgan's avatar
Jeffrey Morgan committed
144
			value:     b.initialValue,
145
146
147
			rate:      0,
			remaining: 0,
		}
Jeffrey Morgan's avatar
Jeffrey Morgan committed
148
	case b.currentValue >= b.maxValue:
149
		b.stats = Stats{
Jeffrey Morgan's avatar
Jeffrey Morgan committed
150
			value:     b.maxValue,
151
152
153
			rate:      0,
			remaining: 0,
		}
154
	default:
155
156
157
158
		rate := b.currentValue - b.stats.value
		var remaining time.Duration
		if rate > 0 {
			remaining = time.Second * time.Duration((float64(b.maxValue-b.currentValue))/(float64(rate)))
159
160
		} else {
			remaining = time.Duration(math.MaxInt64)
161
		}
Michael Yang's avatar
Michael Yang committed
162

163
164
165
166
167
		b.stats = Stats{
			value:     b.currentValue,
			rate:      rate,
			remaining: remaining,
		}
Michael Yang's avatar
Michael Yang committed
168
169
	}

170
171
172
	b.statted = time.Now()

	return b.stats
Michael Yang's avatar
Michael Yang committed
173
}