mirror of
https://github.com/ParisNeo/lollms-webui.git
synced 2024-12-25 06:51:04 +00:00
54 lines
1.5 KiB
Python
54 lines
1.5 KiB
Python
import argparse
|
|
import yaml
|
|
from urllib.parse import urlparse
|
|
|
|
from pathlib import Path
|
|
|
|
def process_yaml(input_file):
|
|
# Read YAML file
|
|
with open(input_file, 'r') as file:
|
|
models = yaml.safe_load(file)
|
|
|
|
# Process each model entry
|
|
for model in models:
|
|
server_url = model['server']
|
|
parsed_url = urlparse(server_url)
|
|
if not 'owner' in model:
|
|
if 'huggingface.co' in parsed_url.netloc:
|
|
# Hugging Face URL, extract owner from server URL
|
|
model['owner'] = parsed_url.path.split('/')[1]
|
|
else:
|
|
# Non-Hugging Face URL, use domain name as owner
|
|
model['owner'] = parsed_url.netloc
|
|
|
|
# Add additional fields
|
|
if not 'link' in model:
|
|
model['link'] = server_url
|
|
if not 'license' in model:
|
|
model['license'] = 'Non commercial'
|
|
|
|
# Save processed YAML file
|
|
output_file = input_file.stem + '_processed.yaml'
|
|
with open(output_file, 'w') as file:
|
|
yaml.dump(models, file)
|
|
|
|
print(f"Processed YAML file saved as {output_file}")
|
|
|
|
def main():
|
|
# Parse command-line arguments
|
|
parser = argparse.ArgumentParser(description='Process YAML file')
|
|
parser.add_argument('input_file', type=str, help='Input YAML file')
|
|
|
|
args = parser.parse_args()
|
|
|
|
input_file = Path(args.input_file)
|
|
|
|
if not input_file.exists():
|
|
print('Input file does not exist.')
|
|
return
|
|
|
|
process_yaml(input_file)
|
|
|
|
if __name__ == '__main__':
|
|
main()
|