bar.go 2.86 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
45
46
47
}

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

func (b *Bar) String() string {
	termWidth, _, err := term.GetSize(int(os.Stderr.Fd()))
	if err != nil {
48
		termWidth = 80
Michael Yang's avatar
Michael Yang committed
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
	}

	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(" ")
	}

67
	fmt.Fprintf(&pre, "%3.0f%% ", math.Floor(b.percent()))
68
	fmt.Fprintf(&suf, "(%s/%s", format.HumanBytes(b.currentValue), format.HumanBytes(b.maxValue))
Michael Yang's avatar
Michael Yang committed
69

70
71
72
73
74
75
76
77
78
79
80
81
82
83
	stats := b.Stats()
	rate := int64(stats.rate)
	if rate > 0 {
		fmt.Fprintf(&suf, ", %s/s", format.HumanBytes(rate))
	}

	fmt.Fprintf(&suf, ")")

	elapsed := time.Since(b.started)
	if b.percent() < 100 && rate > 0 {
		fmt.Fprintf(&suf, " [%s:%s]", elapsed.Round(time.Second), stats.remaining)
	} else {
		fmt.Fprintf(&suf, "        ")
	}
Michael Yang's avatar
Michael Yang committed
84

85
	mid.WriteString("▕")
Michael Yang's avatar
Michael Yang committed
86

87
88
	// 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
89
90
	n := int(float64(f) * b.percent() / 100)

91
92
	if n > 0 {
		mid.WriteString(strings.Repeat("█", n))
Michael Yang's avatar
Michael Yang committed
93
94
95
96
97
98
	}

	if f-n > 0 {
		mid.WriteString(strings.Repeat(" ", f-n))
	}

99
	mid.WriteString("▏")
Michael Yang's avatar
Michael Yang committed
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119

	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
}

120
func (b *Bar) Stats() Stats {
121
	if time.Since(b.statted) < time.Second {
122
		return b.stats
Michael Yang's avatar
Michael Yang committed
123
124
	}

125
126
127
	switch {
	case b.statted.IsZero():
		b.stats = Stats{
Jeffrey Morgan's avatar
Jeffrey Morgan committed
128
			value:     b.initialValue,
129
130
131
			rate:      0,
			remaining: 0,
		}
Jeffrey Morgan's avatar
Jeffrey Morgan committed
132
	case b.currentValue >= b.maxValue:
133
		b.stats = Stats{
Jeffrey Morgan's avatar
Jeffrey Morgan committed
134
			value:     b.maxValue,
135
136
137
			rate:      0,
			remaining: 0,
		}
138
	default:
139
140
141
142
143
		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)))
		}
Michael Yang's avatar
Michael Yang committed
144

145
146
147
148
149
		b.stats = Stats{
			value:     b.currentValue,
			rate:      rate,
			remaining: remaining,
		}
Michael Yang's avatar
Michael Yang committed
150
151
	}

152
153
154
	b.statted = time.Now()

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