fix ModelType()

5ca05c2e · Michael Yang · 7eda70f2 · 5ca05c2e
Commit 5ca05c2e authored Aug 18, 2023 by Michael Yang
Hide whitespace changes
Inline Side-by-side

Showing with 15 additions and 1 deletion

llm/llama.go llm/llama.go +15 -1

No files found.
--- 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 {