Spaces:
Runtime error
Runtime error
File size: 2,559 Bytes
10f3130 307beef 10f3130 |
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 |
import requests
import re
DET_MODELS_FILENAME = 'det_models.txt'
SEG_MODELS_FILENAME = 'seg_models.txt'
CLS_MODELS_FILENAME = 'cls_models.txt'
def load_models_from_txt_files():
"""Load models from txt files."""
with open(DET_MODELS_FILENAME, 'r') as file:
det_models = [line.strip() for line in file]
with open(SEG_MODELS_FILENAME, 'r') as file:
seg_models = [line.strip() for line in file]
with open(CLS_MODELS_FILENAME, 'r') as file:
cls_models = [line.strip() for line in file]
return det_models, seg_models, cls_models
def get_dataset_id_from_model_id(model_id):
"""
Gets the dataset ID from the README file for a given Hugging Face model ID.
Args:
model_id (str): The Hugging Face model ID.
Returns:
The dataset ID as a string, or None if the dataset ID cannot be found.
"""
# Define the URL of the README file for the model
readme_url = f"https://huggingface.co/{model_id}/raw/main/README.md"
# Make a GET request to the README URL and get the contents
response = requests.get(readme_url)
readme_contents = response.text
# Use regular expressions to search for the dataset ID in the README file
match = re.search(r"datasets:\s*\n- (\S+)", readme_contents)
# If a match is found, extract the dataset ID and return it. Otherwise, return None.
if match is not None:
dataset_id = match.group(1)
return dataset_id
else:
return None
def get_task_from_readme(model_id):
"""
Gets the task from the README file for a given Hugging Face model ID.
Args:
model_id (str): The Hugging Face model ID.
Returns:
The task as a string ("detect", "segment", or "classify"), or None if the task cannot be found.
"""
# Define the URL of the README file for the model
readme_url = f"https://huggingface.co/{model_id}/raw/main/README.md"
# Make a GET request to the README URL and get the contents
response = requests.get(readme_url)
readme_contents = response.text
# Use regular expressions to search for the task in the tags section of the README file
if re.search(r"tags:", readme_contents):
if re.search(r"object-detection", readme_contents):
return "detect"
elif re.search(r"image-segmentation", readme_contents):
return "segment"
elif re.search(r"image-classification", readme_contents):
return "classify"
# If the task cannot be found, return None
return None
|