forked from ollama/ollama-python
-
Notifications
You must be signed in to change notification settings - Fork 1
/
model_downloader.py
executable file
·129 lines (107 loc) · 5.11 KB
/
model_downloader.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
#!/usr/bin/env python3
#coding:utf-8
import rospy
import os
import rospkg
import time
import ollama
import tkinter as Tkinter
from subprocess import Popen
class ModelDownloader():
def __init__(self):
Popen(["xterm", "-font", "r16", "-fg", "floralwhite", "-bg", "darkslateblue", "-e", "ollama", "serve"])
time.sleep(1)
self.tk = Tkinter.Tk()
# 以下にモデルの一例を載せる.ここの一例以外のモデルを使いたい場合はREADMEを参照しながら以下のリストに追加してください.
self.can_download_models_info = ["llama3", "llama2", "phi3", "llava", "qwen", "dbrx", "dolphin-mixtral", "llama2-chinese", "tinyllama", "openchat"]
self.download_models_flag = []
self.reset_models_info()
rospack = rospkg.RosPack()
self.iconfile = Tkinter.PhotoImage(file=rospack.get_path("ollama_python")+"/img/icon.png")
self.width = self.tk.winfo_screenwidth()
self.height = self.tk.winfo_screenheight()
self.tk.call('wm', 'iconphoto', self.tk._w, self.iconfile)
def reset_models_info(self):
models = ollama.list()
if (models["models"] is not None):
for m in self.can_download_models_info:
self.download_models_flag += [False]
for lm in models["models"]:
if (m == lm["name"].replace(":latest","")):
self.download_models_flag[-1] = True
break
for lm in models["models"]:
match_flag = False
for m in self.can_download_models_info:
if (m == lm["name"].replace(":latest","")):
match_flag = True
break
if not match_flag:
self.can_download_models_info += [lm["name"].replace(":latest","")]
self.download_models_flag += [True]
else:
for m in self.can_download_models_info:
self.download_models_flag += [False]
def create_gui(self):
# GUI windowの大きさを定義する
geometry_x = 700
if (len(self.can_download_models_info) < 2):
geometry_y = 30 * 2
else:
geometry_y = 30 * len(self.can_download_models_info)
# "window width x window height + position right + position down"
self.tk.geometry(("%sx%s+%s+%s")%(str(geometry_x), str(geometry_y), str((self.width - geometry_x)//2), str((self.height - geometry_y)//2)))
for i, container_info in enumerate(self.can_download_models_info):
if self.download_models_flag[i]:
Tkinter.Button(self.tk, width=9, text="delete" , command=lambda i=i: self.button_clicked_callback("delete", i)).place(x=150, y=i*30)
Tkinter.Button(self.tk, width=9, text="copy" , command=lambda i=i: self.button_clicked_callback("copy" , i)).place(x=250, y=i*30)
Tkinter.Button(self.tk, width=9, text="push" , command=lambda i=i: self.button_clicked_callback("push" , i)).place(x=350, y=i*30)
else:
Tkinter.Button(self.tk, width=34, text="download", command=lambda i=i: self.button_clicked_callback("pull" , i)).place(x=150, y=i*30)
Tkinter.Label(text=container_info, font=("", 15)).place(x=460, y=i*30)
#GUI再起動用のボタンを定義
btn = Tkinter.Button(self.tk, width=6, text="refresh", command=self.refresh_gui)
btn.place(x=0, y=0)
#GUI停止用のボタンを定義
btn = Tkinter.Button(self.tk, width=6, text="close", command=self.quit_gui)
btn.place(x=0, y=30)
self.tk.title("[Download] ollama models GUI")
self.tk.mainloop()
def button_clicked_callback(self, mode, id):
if mode != "dummy":
if mode == "delete":
ollama.delete(str(self.can_download_models_info[id]))
elif mode == "copy":
print("copy")
ollama.copy(str(self.can_download_models_info[id]), str(os.environ.get("USER")) + "/" + str(self.can_download_models_info[id]))
print("copy")
elif mode == "pull":
ollama.pull(str(self.can_download_models_info[id]))
elif mode == "push":
print("push")
ollama.push(str(self.can_download_models_info[id]))
print("push")
self.refresh_gui()
def refresh_gui(self):
self.can_download_models_info = []
self.download_models_flag = []
self.tk.quit()
self.tk.destroy()
self.__init__()
self.create_gui()
def quit_gui(self):
self.tk.quit()
self.tk.destroy()
# def get_window_size():
# # ウィンドウを表示せずに作成する
# # self.tk.withdraw()
# # ウィンドウのサイズを取得する
# width = self.tk.winfo_screenwidth()
# height = self.tk.winfo_screenheight()
# return width, height
def main():
rospy.init_node("model_downloader")
md = ModelDownloader()
md.create_gui()
if __name__ == "__main__":
main()