Spaces:
Running
Running
File size: 1,533 Bytes
df97e68 | 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 | import os
from openai import OpenAI
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()
# Get the API key from environment
api_key = os.getenv("NVIDIA_API_KEY")
if not api_key:
print("Error: NVIDIA_API_KEY not found in .env file.")
# Fallback to check if NVIDIA_API_KEY_2 exists
api_key = os.getenv("NVIDIA_API_KEY_2")
if not api_key:
print("Error: Neither NVIDIA_API_KEY nor NVIDIA_API_KEY_2 found in .env file.")
exit(1)
else:
print("Using NVIDIA_API_KEY_2...")
print(f"Using API Key: {api_key[:10]}...{api_key[-5:]}")
# Initialize the OpenAI client with NVIDIA base URL
client = OpenAI(
base_url="https://integrate.api.nvidia.com/v1",
api_key=api_key
)
print("\nFetching models from NVIDIA API...")
try:
# List models
models = client.models.list()
# Sort models by ID for better readability
sorted_models = sorted(models.data, key=lambda x: x.id)
output_lines = []
output_lines.append(f"{'#':<4} | {'Model ID':<60}")
output_lines.append("-" * 70)
for i, m in enumerate(sorted_models, 1):
output_lines.append(f"{i:<4} | {m.id:<60}")
output_lines.append(f"\nSuccessfully listed {len(models.data)} models.")
# Save to file with UTF-8 encoding
with open("nvidia_models.txt", "w", encoding="utf-8") as f:
f.write("\n".join(output_lines))
# Also print it
print("\n".join(output_lines))
except Exception as e:
print(f"Error accessing NVIDIA API: {e}")
|