baalajimaestro/ollama
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

resume pull download

Browse source
This commit is contained in:
Bruce MacDonald 2023-06-28 15:42:05 -04:00
parent 50f5adb09b
commit b5c54f5a8f
1 changed files with 31 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

25
ollama/model.py
View file

@ -45,30 +45,43 @@ def pull(remote, models_home=".", *args, **kwargs):
json_response = response.json() json_response = response.json()
# get the last bin file we find, this is probably the most up to date
download_url = None
for file_info in json_response: for file_info in json_response:
if file_info.get("type") == "file" and file_info.get("path").endswith(".bin"): if file_info.get("type") == "file" and file_info.get("path").endswith(".bin"):
f_path = file_info.get("path") f_path = file_info.get("path")
download_url = f"https://huggingface.co/{model}/resolve/{branch}/{f_path}" download_url = f"https://huggingface.co/{model}/resolve/{branch}/{f_path}"
if download_url is None:
raise Exception("No model found")
local_filename = os.path.join(models_home, os.path.basename(model)) + ".bin" local_filename = os.path.join(models_home, os.path.basename(model)) + ".bin"
# Check if file already exists
if os.path.exists(local_filename): if os.path.exists(local_filename):
# TODO: check if the file is the same SHA # TODO: check if the file is the same SHA
break first_byte = os.path.getsize(local_filename)
else:
first_byte = 0
response = requests.get(download_url, stream=True) # If file size is non-zero, resume download
if first_byte != 0:
header = {"Range": f"bytes={first_byte}-"}
else:
header = {}
response = requests.get(download_url, headers=header, stream=True)
response.raise_for_status() # Raises stored HTTPError, if one occurred response.raise_for_status() # Raises stored HTTPError, if one occurred
total_size = int(response.headers.get("content-length", 0)) total_size = int(response.headers.get("content-length", 0))
with open(local_filename, "wb") as file, tqdm( with open(local_filename, "ab" if first_byte else "wb") as file, tqdm(
desc=local_filename,
total=total_size, total=total_size,
unit="iB", unit="iB",
unit_scale=True, unit_scale=True,
unit_divisor=1024, unit_divisor=1024,
initial=first_byte,
) as bar: ) as bar:
for data in response.iter_content(chunk_size=1024): for data in response.iter_content(chunk_size=1024):
size = file.write(data) size = file.write(data)
bar.update(size) bar.update(size)
break # Stop after downloading the first .bin file