model_explorer4 / app.py
dwb2023's picture
Update app.py
1e1efc2 verified
raw
history blame
742 Bytes
import gradio as gr
import os
import torch
import spaces
from transformers import AutoModel
from huggingface_hub import login
hf_token = os.getenv("HF_TOKEN")
login(token=hf_token, add_to_git_credential=True)
@spaces.GPU
def get_model_summary(model_name):
# Check if CUDA is available and set the device accordingly
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Load the model and move it to the selected device
model = AutoModel.from_pretrained(model_name, hf_token=hf_token).to(device)
# Return the model's architecture as a string
return str(model)
# Create the Gradio interface
interface = gr.Interface(fn=get_model_summary, inputs="text", outputs="text")
interface.launch()