llm.go 1.7 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374
  1. package llm
  2. import (
  3. "fmt"
  4. "log"
  5. "os"
  6. "github.com/pbnjay/memory"
  7. "github.com/jmorganca/ollama/api"
  8. )
  9. type LLM interface {
  10. Predict([]int, string, func(api.GenerateResponse)) error
  11. Embedding(string) ([]float64, error)
  12. Encode(string) []int
  13. Decode(...int) string
  14. SetOptions(api.Options)
  15. Close()
  16. }
  17. func New(model string, adapters []string, opts api.Options) (LLM, error) {
  18. if _, err := os.Stat(model); err != nil {
  19. return nil, err
  20. }
  21. f, err := os.Open(model)
  22. if err != nil {
  23. return nil, err
  24. }
  25. defer f.Close()
  26. ggml, err := DecodeGGML(f, ModelFamilyLlama)
  27. if err != nil {
  28. return nil, err
  29. }
  30. switch ggml.FileType().String() {
  31. case "F32", "F16", "Q5_0", "Q5_1", "Q8_0":
  32. if opts.NumGPU != 0 {
  33. // F32, F16, Q5_0, Q5_1, and Q8_0 do not support Metal API and will
  34. // cause the runner to segmentation fault so disable GPU
  35. log.Printf("WARNING: GPU disabled for F32, F16, Q5_0, Q5_1, and Q8_0")
  36. opts.NumGPU = 0
  37. }
  38. }
  39. totalResidentMemory := memory.TotalMemory()
  40. switch ggml.ModelType() {
  41. case ModelType3B, ModelType7B:
  42. if totalResidentMemory < 8*1024*1024 {
  43. return nil, fmt.Errorf("model requires at least 8GB of memory")
  44. }
  45. case ModelType13B:
  46. if totalResidentMemory < 16*1024*1024 {
  47. return nil, fmt.Errorf("model requires at least 16GB of memory")
  48. }
  49. case ModelType30B:
  50. if totalResidentMemory < 32*1024*1024 {
  51. return nil, fmt.Errorf("model requires at least 32GB of memory")
  52. }
  53. case ModelType65B:
  54. if totalResidentMemory < 64*1024*1024 {
  55. return nil, fmt.Errorf("model requires at least 64GB of memory")
  56. }
  57. }
  58. switch ggml.ModelFamily() {
  59. case ModelFamilyLlama:
  60. return newLlama(model, adapters, opts)
  61. default:
  62. return nil, fmt.Errorf("unknown ggml type: %s", ggml.ModelFamily())
  63. }
  64. }