Spaces:
Running
Running
File size: 2,196 Bytes
933b7b6 |
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 |
from huggingface_hub import login, model_info, whoami
from transformers import AutoModelForCausalLM, AutoTokenizer
from dotenv import load_dotenv
from datetime import datetime
import json, argparse, sys, os
def is_interactive():
if sys.stdin.isatty():
return True
return hasattr(sys, 'ps1')
def try_interactive_login():
try:
login()
except KeyboardInterrupt:
print("^C")
exit()
def dumps(x):
def jsonable(obj):
d = {}
try:
json.dumps(obj)
except:
if isinstance(obj, datetime):
return obj.isoformat()
try:
d = vars(obj)
except:
return "..."
else:
return obj
for key, value in d.items():
d[key] = jsonable(value)
return d
return json.dumps(jsonable(x), indent=4, separators=(',', ': '))
parser = argparse.ArgumentParser()
parser.add_argument('model_id', type=str)
args = parser.parse_args()
model_id = args.model_id
path = os.path.join(os.getcwd(), ".hf_home", model_id)
print("Downloading to", path)
if os.path.exists(path):
print(f"{path} already exists, aborting. (To redownload, rm it first).")
exit()
try:
model_info(model_id)
except Exception as e:
print(e)
exit(1)
try:
user_info = whoami()
except Exception as e:
huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
if not huggingface_token:
load_dotenv()
huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
if huggingface_token:
print("Logging in with env HUGGINGFACE_TOKEN")
try:
login(huggingface_token)
except Exception as e:
print(e)
try_interactive_login()
elif is_interactive():
try_interactive_login()
else:
print("Missing env: HUGGINGFACE_TOKEN")
exit(1)
user_info = whoami()
print("Authenticated as:", dumps(user_info))
try:
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)
except OSError as e:
print(e)
exit(1)
tokenizer.save_pretrained(path)
model.save_pretrained(path)
|