mixtral.go 2.4 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788
  1. package convert
  2. import (
  3. "io"
  4. "regexp"
  5. "github.com/ollama/ollama/llm"
  6. )
  7. type MixtralModel struct {
  8. ModelData
  9. }
  10. func (m *MixtralModel) GetTensors() error {
  11. t, err := m.Format.GetTensors(m.Path, m.Params)
  12. if err != nil {
  13. return err
  14. }
  15. pattern := `^blk\.[0-9]+\.attn_(?P<layer>q|k)\.weight$`
  16. re, err := regexp.Compile(pattern)
  17. if err != nil {
  18. return err
  19. }
  20. for _, l := range t {
  21. matches := re.FindAllStringSubmatch(l.Name, -1)
  22. if len(matches) > 0 {
  23. wt := l.WriterTo.(safetensorWriterTo)
  24. wt.repacker = m.Repack
  25. l.WriterTo = wt
  26. }
  27. m.Tensors = append(m.Tensors, l)
  28. }
  29. return nil
  30. }
  31. func (m *MixtralModel) LoadVocab() error {
  32. v, err := LoadSentencePieceTokens(m.Path, m.Params)
  33. if err != nil {
  34. return err
  35. }
  36. m.Vocab = v
  37. return nil
  38. }
  39. func (m *MixtralModel) WriteGGUF(ws io.WriteSeeker) error {
  40. kv := llm.KV{
  41. "general.architecture": "llama",
  42. "general.name": m.Name,
  43. "llama.block_count": uint32(m.Params.HiddenLayers),
  44. "llama.context_length": uint32(m.Params.ContextSize),
  45. "llama.embedding_length": uint32(m.Params.HiddenSize),
  46. "llama.feed_forward_length": uint32(m.Params.IntermediateSize),
  47. "llama.attention.head_count": uint32(m.Params.AttentionHeads),
  48. "llama.attention.head_count_kv": uint32(m.Params.KeyValHeads),
  49. "llama.rope.freq_base": float32(m.Params.RopeFrequencyBase),
  50. "llama.attention.layer_norm_rms_epsilon": float32(m.Params.NormEPS),
  51. "llama.expert_count": uint32(m.Params.Experts),
  52. "llama.expert_used_count": uint32(m.Params.ExpertsUsed),
  53. "llama.vocab_size": uint32(len(m.Vocab.Tokens)),
  54. "llama.rope.dimension_count": uint32(m.Params.HiddenSize / m.Params.AttentionHeads),
  55. "general.file_type": uint32(1),
  56. "tokenizer.ggml.model": "llama",
  57. "tokenizer.ggml.tokens": m.Vocab.Tokens,
  58. "tokenizer.ggml.scores": m.Vocab.Scores,
  59. "tokenizer.ggml.token_type": m.Vocab.Types,
  60. "tokenizer.ggml.bos_token_id": uint32(m.Params.BoSTokenID),
  61. "tokenizer.ggml.eos_token_id": uint32(m.Params.EoSTokenID),
  62. "tokenizer.ggml.unknown_token_id": uint32(0),
  63. "tokenizer.ggml.add_bos_token": true,
  64. "tokenizer.ggml.add_eos_token": false,
  65. }
  66. return llm.NewGGUFV3(m.Params.ByteOrder).Encode(ws, kv, m.Tensors)
  67. }
  68. func (m *MixtralModel) Repack(name string, data []float32, shape []uint64) ([]float32, error) {
  69. return llamaRepack(name, m.Params, data, shape)
  70. }