This commit is contained in:
Josh Yan 2024-07-16 11:21:13 -07:00
parent 25be20949c
commit fc1b3ee9bf
2 changed files with 45 additions and 35 deletions

View File

@ -15,7 +15,7 @@ import (
func TestGGUFRewrite(t *testing.T) {
tests := []string{
"glm2.gguf",
"phi3.gguf",
}
for i := range tests {
@ -34,7 +34,7 @@ func TestGGUFRewrite(t *testing.T) {
}
defer f.Close()
ggml, m, err := decodeGGML(t, f)
ggml, _, err := decodeGGML(t, f)
if err != nil {
t.Fatal(err)
}
@ -45,11 +45,11 @@ func TestGGUFRewrite(t *testing.T) {
}
defer temp.Close()
n, ggml2, err := rewriteGGML(t, ggml, temp)
n, ggml2, err := rewriteGGML(t, ggml, temp, f)
if n != m {
/* if n != m {
t.Fatalf("n: %d, m: %d", n, m)
}
} */
if err != nil {
t.Fatal(err)
@ -57,31 +57,9 @@ func TestGGUFRewrite(t *testing.T) {
if diff, diff2, ok := compareGGML(n, ggml2, ggml, temp, f); !ok {
if cmp.Diff(diff, diff2) != "" {
t.Fatalf("\n%s,\n%s\ndiff: %s", diff["token_embd.weight"], diff2["token_embd.weight"], cmp.Diff(diff, diff2))
t.Fatalf("\n%s,\n%s\n%s\n%s\ndiff: %s", diff["token_embd.weight"], diff2["token_embd.weight"], diff["token_embd.weight size"], diff["token_embd.weight offset"], cmp.Diff(diff, diff2))
}
}
/* // Reset the file offset to the beginning
if _, err := f.Seek(0, io.SeekStart); err != nil {
t.Fatal(err)
}
if _, err := temp.Seek(0, io.SeekStart); err != nil {
t.Fatal(err)
}
content1, err := io.ReadAll(f)
if err != nil {
t.Fatalf("failed to read file1: %v", err)
}
content2, err := io.ReadAll(temp)
if err != nil {
t.Fatalf("failed to read file1: %v", err)
}
if byteCmp := cmp.Diff(content1, content2); byteCmp != "" {
t.Fatalf("content diff: %s", byteCmp)
} */
})
}
}
@ -121,6 +99,10 @@ func compareGGML(n int64, ggml1, ggml2 *GGML, f *os.File, f2 *os.File) (map[stri
if diffy := cmp.Diff(t.values, kv2[k].(*array).values); diffy != "" {
diff[k] = diffy
}
default:
if v != kv2[k] {
diff[k] = fmt.Sprintf("kv1: %v, kv2: %v", v, kv2[k])
}
}
// }
@ -136,21 +118,46 @@ func compareGGML(n int64, ggml1, ggml2 *GGML, f *os.File, f2 *os.File) (map[stri
for _, tensor := range t1 {
sha256sum := sha256.New()
sr := io.NewSectionReader(f, n+int64(tensor.Offset), int64(tensor.Size()))
if _, err := io.Copy(sha256sum, sr); err != nil {
var s int64
s, err := io.Copy(sha256sum, sr)
if err != nil {
fmt.Println(err)
}
diff[tensor.Name] = fmt.Sprintf("%x", sha256sum.Sum(nil))
diff[tensor.Name+" size"] = fmt.Sprintf("%d", s)
diff[tensor.Name+" offset"] = fmt.Sprintf("%v", tensor.Offset)
}
/* sha256Sum2 := sha256.New()
sr1 := io.NewSectionReader(f2, 0, n)
s1, err := io.Copy(sha256Sum2, sr1)
if err != nil {
return nil, nil, true
}
sha256Sum3 := sha256.New()
sr2 := io.NewSectionReader(f, 0, n)
s2, err := io.Copy(sha256Sum3, sr2)
if err != nil {
return nil, nil, true
}
diff["sha"] = fmt.Sprintf("%d", s1)
diff2["sha"] = fmt.Sprintf("%d", s2) */
for _, tensor := range t2 {
sha256sum := sha256.New()
var s int64
sr := io.NewSectionReader(f2, n+int64(tensor.Offset), int64(tensor.Size()))
if _, err := io.Copy(sha256sum, sr); err != nil {
s, err := io.Copy(sha256sum, sr)
if err != nil {
fmt.Println(err)
}
diff2[tensor.Name] = fmt.Sprintf("%x", sha256sum.Sum(nil))
diff2[tensor.Name+" size"] = fmt.Sprintf("%d", s)
diff2[tensor.Name+" offset"] = fmt.Sprintf("%v", tensor.Offset)
}
return diff, diff2, len(diff) == 0
@ -164,22 +171,24 @@ func decodeGGML(t *testing.T, f *os.File) (*GGML, int64, error) {
return ggml, n, nil
}
func rewriteGGML(t *testing.T, ggml *GGML, temp *os.File) (int64, *GGML, error) {
func rewriteGGML(t *testing.T, ggml *GGML, temp *os.File, f *os.File) (int64, *GGML, error) {
var tensors Tensors
fmt.Println("11111111111111111111111111111111111111111")
for _, tensor := range ggml.Tensors() {
shape := make([]uint64, len(tensor.Shape))
for i := range len(tensor.Shape) {
shape[i] = tensor.Shape[len(tensor.Shape)-i-1]
}
fmt.Println("tensors", tensor.Name, shape, tensor.Kind, 737414+int64(tensor.Offset))
tensors = append(tensors, &Tensor{
Name: tensor.Name,
Kind: tensor.Kind,
Shape: shape,
WriterTo: TensorWriter{
Reader: io.NewSectionReader(temp, int64(tensor.Offset), int64(tensor.Size())),
Reader: io.NewSectionReader(f, 737414+int64(tensor.Offset), int64(tensor.Size())),
},
})
}
@ -195,13 +204,13 @@ func rewriteGGML(t *testing.T, ggml *GGML, temp *os.File) (int64, *GGML, error)
t.Fatal(err)
}
fmt.Println(n)
temp.Seek(0, io.SeekStart)
fmt.Println(n, " is my offset")
file, err := os.Open(temp.Name())
if err != nil {
t.Fatal(err)
}
ggml2, n, err := DecodeGGML(file, math.MaxInt)
ggml2, _, err := DecodeGGML(file, math.MaxInt)
if err != nil {
t.Fatal(err)
}

View File

@ -259,6 +259,7 @@ func parseFromFile(ctx context.Context, file *os.File, digest string, fn func(ap
Shape: shape,
WriterTo: &llm.TensorWriter{
// This needs offset + tensors.Offset int64(tensor.Offset) to be correct
Reader: io.NewSectionReader(file, int64(tensor.Offset), int64(tensor.Size())),
},
})