From 92189a58553ac88edf19ea5a8834478bd5fd6d86 Mon Sep 17 00:00:00 2001
From: Michael Yang <mxyng@pm.me>
Date: Thu, 12 Oct 2023 09:34:16 -0700
Subject: [PATCH] fix memory check

---
 api/client.go   |  3 ++-
 format/bytes.go | 19 +++++++++++++------
 llm/llama.go    |  2 +-
 llm/llm.go      | 21 +++++++++++----------
 4 files changed, 27 insertions(+), 18 deletions(-)

diff --git a/api/client.go b/api/client.go
index f308b233..14ea353d 100644
--- a/api/client.go
+++ b/api/client.go
@@ -14,6 +14,7 @@ import (
 	"runtime"
 	"strings"
 
+	"github.com/jmorganca/ollama/format"
 	"github.com/jmorganca/ollama/version"
 )
 
@@ -127,7 +128,7 @@ func (c *Client) do(ctx context.Context, method, path string, reqData, respData
 	return nil
 }
 
-const maxBufferSize = 512 * 1000 // 512KB
+const maxBufferSize = 512 * format.KiloByte
 
 func (c *Client) stream(ctx context.Context, method, path string, data any, fn func([]byte) error) error {
 	var buf *bytes.Buffer
diff --git a/format/bytes.go b/format/bytes.go
index 63cc7b00..ca5ac640 100644
--- a/format/bytes.go
+++ b/format/bytes.go
@@ -2,14 +2,21 @@ package format
 
 import "fmt"
 
+const (
+	Byte     = 1
+	KiloByte = Byte * 1000
+	MegaByte = KiloByte * 1000
+	GigaByte = MegaByte * 1000
+)
+
 func HumanBytes(b int64) string {
 	switch {
-	case b > 1000*1000*1000:
-		return fmt.Sprintf("%d GB", b/1000/1000/1000)
-	case b > 1000*1000:
-		return fmt.Sprintf("%d MB", b/1000/1000)
-	case b > 1000:
-		return fmt.Sprintf("%d KB", b/1000)
+	case b > GigaByte:
+		return fmt.Sprintf("%d GB", b/GigaByte)
+	case b > MegaByte:
+		return fmt.Sprintf("%d MB", b/MegaByte)
+	case b > KiloByte:
+		return fmt.Sprintf("%d KB", b/KiloByte)
 	default:
 		return fmt.Sprintf("%d B", b)
 	}
diff --git a/llm/llama.go b/llm/llama.go
index 0b460e9a..db51429c 100644
--- a/llm/llama.go
+++ b/llm/llama.go
@@ -509,7 +509,7 @@ type PredictRequest struct {
 	Stop             []string `json:"stop,omitempty"`
 }
 
-const maxBufferSize = 512 * 1000 // 512KB
+const maxBufferSize = 512 * format.KiloByte
 
 func (llm *llama) Predict(ctx context.Context, prevContext []int, prompt string, fn func(api.GenerateResponse)) error {
 	prevConvo, err := llm.Decode(ctx, prevContext)
diff --git a/llm/llm.go b/llm/llm.go
index 6df2a47c..4ae5dd2e 100644
--- a/llm/llm.go
+++ b/llm/llm.go
@@ -10,6 +10,7 @@ import (
 	"github.com/pbnjay/memory"
 
 	"github.com/jmorganca/ollama/api"
+	"github.com/jmorganca/ollama/format"
 )
 
 type LLM interface {
@@ -60,33 +61,33 @@ func New(workDir, model string, adapters []string, opts api.Options) (LLM, error
 	totalResidentMemory := memory.TotalMemory()
 	switch ggml.ModelType() {
 	case "3B", "7B":
-		if ggml.FileType() == "F16" && totalResidentMemory < 16*1000*1000 {
+		if ggml.FileType() == "F16" && totalResidentMemory < 16*format.GigaByte {
 			return nil, fmt.Errorf("F16 model requires at least 16 GB of memory")
-		} else if totalResidentMemory < 8*1000*1000 {
+		} else if totalResidentMemory < 8*format.GigaByte {
 			return nil, fmt.Errorf("model requires at least 8 GB of memory")
 		}
 	case "13B":
-		if ggml.FileType() == "F16" && totalResidentMemory < 32*1000*1000 {
+		if ggml.FileType() == "F16" && totalResidentMemory < 32*format.GigaByte {
 			return nil, fmt.Errorf("F16 model requires at least 32 GB of memory")
-		} else if totalResidentMemory < 16*1000*1000 {
+		} else if totalResidentMemory < 16*format.GigaByte {
 			return nil, fmt.Errorf("model requires at least 16 GB of memory")
 		}
 	case "30B", "34B", "40B":
-		if ggml.FileType() == "F16" && totalResidentMemory < 64*1000*1000 {
+		if ggml.FileType() == "F16" && totalResidentMemory < 64*format.GigaByte {
 			return nil, fmt.Errorf("F16 model requires at least 64 GB of memory")
-		} else if totalResidentMemory < 32*1000*1000 {
+		} else if totalResidentMemory < 32*format.GigaByte {
 			return nil, fmt.Errorf("model requires at least 32 GB of memory")
 		}
 	case "65B", "70B":
-		if ggml.FileType() == "F16" && totalResidentMemory < 128*1000*1000 {
+		if ggml.FileType() == "F16" && totalResidentMemory < 128*format.GigaByte {
 			return nil, fmt.Errorf("F16 model requires at least 128 GB of memory")
-		} else if totalResidentMemory < 64*1000*1000 {
+		} else if totalResidentMemory < 64*format.GigaByte {
 			return nil, fmt.Errorf("model requires at least 64 GB of memory")
 		}
 	case "180B":
-		if ggml.FileType() == "F16" && totalResidentMemory < 512*1000*1000 {
+		if ggml.FileType() == "F16" && totalResidentMemory < 512*format.GigaByte {
 			return nil, fmt.Errorf("F16 model requires at least 512GB of memory")
-		} else if totalResidentMemory < 128*1000*1000 {
+		} else if totalResidentMemory < 128*format.GigaByte {
 			return nil, fmt.Errorf("model requires at least 128GB of memory")
 		}
 	}