-
Notifications
You must be signed in to change notification settings - Fork 4
/
utils.py
59 lines (47 loc) · 1.8 KB
/
utils.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
def get_pdf_image(doc_path):
"""Converts a PDF document to an image."""
from wand.image import Image as WImage
from pathlib import Path
img = None
if Path(doc_path).exists() and Path(doc_path).is_file():
img = WImage(filename=doc_path)
# make background of img white
img.format = "png"
img.background_color = "white"
img.alpha_channel = "remove"
return img
def update_state(key, value):
"""Update the .env.state state file with the key and value."""
from pathlib import Path
from dotenv import dotenv_values
data = {}
state_filename = ".env.state"
if Path(state_filename).exists() and Path(state_filename).is_file():
data = dotenv_values(state_filename)
data[key] = value
print(f"Updating state file with {key}={redact_secret(key, value)}")
with open(state_filename, "w") as f:
for k, v in data.items():
f.write(f"{k}={v}\n")
def redact_secret(key, value):
"""Redact a value from the logs if the key indicates that the value contains a keyword such as KEY or SECRET."""
if "KEY" in key or "SECRET" in key:
return value[:4] + "*" * (len(value) - 4)
return value
def wait_for_model(client, model_name):
"""Wait for the model to be available, typically waiting for a finetuning job to complete."""
import time
attempts = 0
while True:
try:
model = client.models.get(model_name, label="latest")
return model
except:
print(f"Model not found yet #{attempts}")
attempts += 1
time.sleep(30)
def file_sha256(filename):
"""Returns the SHA256 hash of a file."""
import hashlib
with open(filename, "rb", buffering=0) as f:
return hashlib.file_digest(f, "sha256").hexdigest()