浏览代码

use cut prefix

Michael Yang 1 年之前
父节点
当前提交
b36b0b71f8
共有 1 个文件被更改,包括 4 次插入6 次删除
  1. 4 6
      llm/llama.go

+ 4 - 6
llm/llama.go

@@ -531,16 +531,14 @@ func (llm *llama) Predict(ctx context.Context, prevContext []int, prompt string,
 			// This handles the request cancellation
 			return ctx.Err()
 		default:
-			line := scanner.Text()
-			if line == "" {
+			line := scanner.Bytes()
+			if len(line) == 0 {
 				continue
 			}
 
-			// Read data from the server-side event stream
-			if strings.HasPrefix(line, "data: ") {
-				evt := line[6:]
+			if evt, ok := bytes.CutPrefix(line, []byte("data: ")); ok {
 				var p prediction
-				if err := json.Unmarshal([]byte(evt), &p); err != nil {
+				if err := json.Unmarshal(evt, &p); err != nil {
 					return fmt.Errorf("error unmarshaling llm prediction response: %v", err)
 				}