model_explorer2 / app.py
dwb2023's picture
Update app.py
028d122 verified
raw
history blame
1.42 kB
import gradio as gr
import os
import torch, torchvision, einops
import spaces
from transformers import AutoModelForCausalLM
from huggingface_hub import login
import subprocess
subprocess.run(
"pip install flash-attn --no-build-isolation",
env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
shell=True,
)
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 = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True).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",
examples=[
["google/gemma-7b"],
["microsoft/Phi-3-mini-4k-instruct"],
["meta-llama/Meta-Llama-3-8B"],
["mistralai/Mistral-7B-Instruct-v0.3"],
["vikhyatk/moondream2"],
["microsoft/Phi-3-vision-128k-instruct"],
["openbmb/MiniCPM-Llama3-V-2_5"],
["google/paligemma-3b-mix-224"],
["HuggingFaceM4/idefics2-8b-chatty"],
["mistralai/Codestral-22B-v0.1"]
]
)
interface.launch()