From 5ca05c2e88c8d605b33bafc241fc84b77d322fc4 Mon Sep 17 00:00:00 2001
From: Michael Yang <mxyng@pm.me>
Date: Fri, 18 Aug 2023 11:23:31 -0700
Subject: [PATCH] fix ModelType()

---
 llm/llama.go | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/llm/llama.go b/llm/llama.go
index 8c5762b6..ce697b33 100644
--- a/llm/llama.go
+++ b/llm/llama.go
@@ -117,7 +117,21 @@ func (llm *llamaModel) ModelFamily() ModelFamily {
 }
 
 func (llm *llamaModel) ModelType() ModelType {
-	return ModelType30B
+	switch llm.hyperparameters.NumLayer {
+	case 26:
+		return ModelType3B
+	case 32:
+		return ModelType7B
+	case 40:
+		return ModelType13B
+	case 60:
+		return ModelType30B
+	case 80:
+		return ModelType65B
+	}
+
+	// TODO: find a better default
+	return ModelType7B
 }
 
 func (llm *llamaModel) FileType() FileType {