1
0
mirror of https://github.com/fumiama/deepinfra.git synced 2026-06-05 00:32:46 +08:00

feat(model): add item pre/post modelize hook

This commit is contained in:
源文雨
2025-09-24 21:41:53 +08:00
parent 8c92983b13
commit 348633b5c6
2 changed files with 22 additions and 2 deletions

View File

@@ -22,9 +22,16 @@ func (l *Log[T]) newbatch(sz int) *batch[T] {
func (cl *batch[T]) String() string { func (cl *batch[T]) String() string {
sb := strings.Builder{} sb := strings.Builder{}
for _, item := range cl.items { for i := range cl.items {
pre, post := cl.lst.preM, cl.lst.postM
if pre != nil {
pre(&cl.items[i])
}
sb.WriteString(cl.lst.itemsep) sb.WriteString(cl.lst.itemsep)
sb.WriteString(item.String()) sb.WriteString(cl.items[i].String())
if post != nil {
post(&cl.items[i])
}
} }
return sb.String()[len(cl.lst.itemsep):] return sb.String()[len(cl.lst.itemsep):]
} }

View File

@@ -14,6 +14,7 @@ type Log[T fmt.Stringer] struct {
itemsep string itemsep string
defaultprompt string defaultprompt string
m map[int64][]*batch[T] m map[int64][]*batch[T]
preM, postM func(*T)
} }
func NewLog[T fmt.Stringer](batchcap, itemscap int, itemsep, defaultprompt string) Log[T] { func NewLog[T fmt.Stringer](batchcap, itemscap int, itemsep, defaultprompt string) Log[T] {
@@ -35,6 +36,18 @@ func NewLog[T fmt.Stringer](batchcap, itemscap int, itemsep, defaultprompt strin
} }
} }
// SetPreModelize will be called before calling item.String()
func (l *Log[T]) SetPreModelize(fn func(*T)) *Log[T] {
l.preM = fn
return l
}
// SetPreModelize will be called after calling item.String()
func (l *Log[T]) SetPostModelize(fn func(*T)) *Log[T] {
l.postM = fn
return l
}
func (l *Log[T]) Add(grp int64, item T, isbot bool) { func (l *Log[T]) Add(grp int64, item T, isbot bool) {
l.mu.Lock() l.mu.Lock()
defer l.mu.Unlock() defer l.mu.Unlock()