Spaces:
Paused
Paused
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) | |
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() | |