load-balancer / hf_scrapper.py
ChandimaPrabath's picture
0.0.0.1 Alpha
836f75b
import os
import requests
import json
from requests.exceptions import RequestException
def get_file_structure(repo, token, path=""):
"""
Fetches the file structure of a specified Hugging Face repository.
Args:
repo (str): The name of the repository.
token (str): The authorization token for the request.
path (str, optional): The specific path in the repository. Defaults to "".
Returns:
list: A list of file structure information.
"""
api_url = f"https://huggingface.co/api/models/{repo}/tree/main/{path}"
headers = {'Authorization': f'Bearer {token}'}
print(f"Fetching file structure from URL: {api_url}")
try:
response = requests.get(api_url, headers=headers)
response.raise_for_status()
return response.json()
except RequestException as e:
print(f"Error fetching file structure: {e}")
return []
def write_file_structure_to_json(file_structure, file_path):
"""
Writes the file structure to a JSON file.
Args:
file_structure (list): The file structure data.
file_path (str): The path where the JSON file will be saved.
"""
try:
with open(file_path, 'w') as json_file:
json.dump(file_structure, json_file, indent=2)
print(f'File structure written to {file_path}')
except IOError as e:
print(f"Error writing file structure to JSON: {e}")