-
-
Save mseri/fae3d72ddd9d117b8b937c94c6496a28 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