Created
June 14, 2024 22:47
-
-
Save JerrettDavis/7bc86098e705e3a7b4efcd60a2b413d7 to your computer and use it in GitHub Desktop.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| package main | |
| import ( | |
| "encoding/json" | |
| "fmt" | |
| "io/ioutil" | |
| "os" | |
| "os/user" | |
| "path/filepath" | |
| "strings" | |
| ) | |
| const ( | |
| successPrefix = "\033[1;32mSuccess:\033[0m" | |
| failedPrefix = "\033[0;31mFailed:\033[0m" | |
| ) | |
| type Layer struct { | |
| Digest string `json:"digest"` | |
| MediaType string `json:"mediaType"` | |
| } | |
| type Manifest struct { | |
| Layers []Layer `json:"layers"` | |
| } | |
| func main() { | |
| if len(os.Args) < 3 { | |
| fmt.Println("Usage: go run main.go <model_name> <target_path>") | |
| return | |
| } | |
| modelName := os.Args[1] | |
| targetPath := os.Args[2] | |
| usr, err := user.Current() | |
| if err != nil { | |
| fmt.Printf("%s Failed to get current user: %v\n", failedPrefix, err) | |
| return | |
| } | |
| ollamaHome := filepath.Join(usr.HomeDir, ".ollama") | |
| blobsFileBasePath := filepath.Join(ollamaHome, "models", "blobs") | |
| manifestsFileBasePath := filepath.Join(ollamaHome, "models", "manifests") | |
| nameArgs := strings.Split(strings.ReplaceAll(modelName, ":", "/"), "/") | |
| var manifestsRegistryName, manifestsLibraryName, manifestsModelName, manifestsParamsName string | |
| switch len(nameArgs) { | |
| case 4: | |
| manifestsRegistryName = nameArgs[0] | |
| manifestsLibraryName = nameArgs[1] | |
| manifestsModelName = nameArgs[2] | |
| manifestsParamsName = nameArgs[3] | |
| case 3: | |
| manifestsLibraryName = nameArgs[0] | |
| manifestsModelName = nameArgs[1] | |
| manifestsParamsName = nameArgs[2] | |
| case 2: | |
| manifestsModelName = nameArgs[0] | |
| manifestsParamsName = nameArgs[1] | |
| case 1: | |
| manifestsModelName = nameArgs[0] | |
| } | |
| if manifestsRegistryName == "" { | |
| manifestsRegistryName = "registry.ollama.ai" | |
| } | |
| if manifestsLibraryName == "" { | |
| manifestsLibraryName = "library" | |
| } | |
| if manifestsModelName == "" { | |
| manifestsModelName = "vicuna" | |
| } | |
| if manifestsParamsName == "" { | |
| manifestsParamsName = "latest" | |
| } | |
| modelFullName := manifestsModelName + ":" + manifestsParamsName | |
| fmt.Printf("Exporting model \"%s\" to \"%s\"...\n\n", modelFullName, targetPath) | |
| manifestsFilePath := filepath.Join(manifestsFileBasePath, manifestsRegistryName, manifestsLibraryName, manifestsModelName, manifestsParamsName) | |
| if _, err := os.Stat(manifestsFilePath); os.IsNotExist(err) { | |
| fmt.Printf("%s \"%s\" does not exist, the model \"%s\" you requested is not found.\n", failedPrefix, manifestsFilePath, modelFullName) | |
| return | |
| } | |
| if _, err := os.Stat(targetPath); err == nil { | |
| fmt.Printf("%s \"%s\" already exists, exiting to prevent unexpected operations.\n", failedPrefix, targetPath) | |
| return | |
| } | |
| os.MkdirAll(targetPath, os.ModePerm) | |
| sourceFilePath := filepath.Join(targetPath, "source.txt") | |
| ioutil.WriteFile(sourceFilePath, []byte(fmt.Sprintf("%s/%s/%s:%s", manifestsRegistryName, manifestsLibraryName, manifestsModelName, manifestsParamsName)), os.ModePerm) | |
| manifestData, err := ioutil.ReadFile(manifestsFilePath) | |
| if err != nil { | |
| fmt.Printf("%s Failed to read manifest file: %v\n", failedPrefix, err) | |
| return | |
| } | |
| var manifest Manifest | |
| err = json.Unmarshal(manifestData, &manifest) | |
| if err != nil { | |
| fmt.Printf("%s Failed to unmarshal manifest data: %v\n", failedPrefix, err) | |
| return | |
| } | |
| exportModelFilePath := filepath.Join(targetPath, "Modelfile") | |
| exportModelBinPath := filepath.Join(targetPath, "model.bin") | |
| for _, layer := range manifest.Layers { | |
| blobFileName := strings.ReplaceAll(layer.Digest, ":", "-") | |
| blobFilePath := filepath.Join(blobsFileBasePath, blobFileName) | |
| blobData, err := ioutil.ReadFile(blobFilePath) | |
| if err != nil { | |
| fmt.Printf("%s Failed to read blob file: %v\n", failedPrefix, err) | |
| return | |
| } | |
| blobTypeName := strings.Split(layer.MediaType, ".")[len(strings.Split(layer.MediaType, "."))-1] | |
| switch blobTypeName { | |
| case "model": | |
| ioutil.WriteFile(exportModelBinPath, blobData, os.ModePerm) | |
| appendToFile(exportModelFilePath, "FROM ./model.bin\n") | |
| case "params": | |
| paramsJson := string(blobData) | |
| paramsMap := make(map[string]interface{}) | |
| json.Unmarshal([]byte(paramsJson), ¶msMap) | |
| for key, value := range paramsMap { | |
| switch v := value.(type) { | |
| case []interface{}: | |
| for _, val := range v { | |
| appendToFile(exportModelFilePath, fmt.Sprintf("PARAMETER %s \"%v\"\n", key, val)) | |
| } | |
| } | |
| } | |
| default: | |
| typeName := strings.ToUpper(blobTypeName) | |
| appendToFile(exportModelFilePath, fmt.Sprintf("%s \"\"\"%s\"\"\"\n", typeName, string(blobData))) | |
| } | |
| } | |
| fmt.Printf("%s Model \"%s\" has been exported to \"%s\"!\n", successPrefix, modelFullName, targetPath) | |
| } | |
| func appendToFile(filePath, text string) { | |
| f, err := os.OpenFile(filePath, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0600) | |
| if err != nil { | |
| fmt.Printf("%s Failed to open file for appending: %v\n", failedPrefix, err) | |
| return | |
| } | |
| defer f.Close() | |
| if _, err = f.WriteString(text); err != nil { | |
| fmt.Printf("%s Failed to write to file: %v\n", failedPrefix, err) | |
| } | |
| } |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment