bar.go 3.41 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
	suf.WriteString(strings.Repeat(" ", 44-suf.Len()-len(timing)))
99
	suf.WriteString(timing)
Michael Yang's avatar
Michael Yang committed
100

101
102
	// 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
103
104
	n := int(float64(f) * b.percent() / 100)

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

	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
}

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

138
139
140
	switch {
	case b.statted.IsZero():
		b.stats = Stats{
Jeffrey Morgan's avatar
Jeffrey Morgan committed
141
			value:     b.initialValue,
142
143
144
			rate:      0,
			remaining: 0,
		}
Jeffrey Morgan's avatar
Jeffrey Morgan committed
145
	case b.currentValue >= b.maxValue:
146
		b.stats = Stats{
Jeffrey Morgan's avatar
Jeffrey Morgan committed
147
			value:     b.maxValue,
148
149
150
			rate:      0,
			remaining: 0,
		}
151
	default:
152
153
154
155
		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)))
156
157
		} else {
			remaining = time.Duration(math.MaxInt64)
158
		}
Michael Yang's avatar
Michael Yang committed
159

160
161
162
163
164
		b.stats = Stats{
			value:     b.currentValue,
			rate:      rate,
			remaining: remaining,
		}
Michael Yang's avatar
Michael Yang committed
165
166
	}

167
168
169
	b.statted = time.Now()

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