Spaces:
Runtime error
Runtime error
| import logging | |
| import os | |
| import subprocess | |
| from typing import Callable, Dict, List | |
| import psutil | |
| import torch | |
| logger = logging.getLogger(__name__) | |
| def log_available_gpu_memory(): | |
| """ | |
| Logs the available GPU memory for each available GPU device. | |
| If no GPUs are available, logs "No GPUs available". | |
| Returns: | |
| None | |
| """ | |
| if torch.cuda.is_available(): | |
| for i in range(torch.cuda.device_count()): | |
| memory_info = subprocess.check_output( | |
| f"nvidia-smi -i {i} --query-gpu=memory.free --format=csv,nounits,noheader", | |
| shell=True, | |
| ) | |
| memory_info = str(memory_info).split("\\")[0][2:] | |
| logger.info(f"GPU {i} memory available: {memory_info} MiB") | |
| else: | |
| logger.info("No GPUs available") | |
| def log_available_ram(): | |
| """ | |
| Logs the amount of available RAM in gigabytes. | |
| Returns: | |
| None | |
| """ | |
| memory_info = psutil.virtual_memory() | |
| # convert bytes to GB | |
| logger.info(f"Available RAM: {memory_info.available / (1024.0 ** 3):.2f} GB") | |
| def log_files_and_subdirs(directory_path: str): | |
| """ | |
| Logs all files and subdirectories in the specified directory. | |
| Args: | |
| directory_path (str): The path to the directory to log. | |
| Returns: | |
| None | |
| """ | |
| # Check if the directory exists | |
| if os.path.exists(directory_path) and os.path.isdir(directory_path): | |
| for dirpath, dirnames, filenames in os.walk(directory_path): | |
| logger.info(f"Directory: {dirpath}") | |
| for filename in filenames: | |
| logger.info(f"File: {os.path.join(dirpath, filename)}") | |
| for dirname in dirnames: | |
| logger.info(f"Sub-directory: {os.path.join(dirpath, dirname)}") | |
| else: | |
| logger.info(f"The directory '{directory_path}' does not exist") | |
| class MockedPipeline: | |
| """ | |
| A mocked pipeline class that is used as a replacement to the HF pipeline class. | |
| Attributes: | |
| ----------- | |
| task : str | |
| The task of the pipeline, which is text-generation. | |
| f : Callable[[str], str] | |
| A function that takes a prompt string as input and returns a generated text string. | |
| """ | |
| task: str = "text-generation" | |
| def __init__(self, f: Callable[[str], str]): | |
| self.f = f | |
| def __call__(self, prompt: str) -> List[Dict[str, str]]: | |
| """ | |
| Calls the pipeline with a given prompt and returns a list of generated text. | |
| Parameters: | |
| ----------- | |
| prompt : str | |
| The prompt string to generate text from. | |
| Returns: | |
| -------- | |
| List[Dict[str, str]] | |
| A list of dictionaries, where each dictionary contains a generated_text key with the generated text string. | |
| """ | |
| result = self.f(prompt) | |
| return [{"generated_text": f"{prompt}{result}"}] | |