Spaces:
Running
Running
File size: 6,475 Bytes
ec1b552 92bf0a5 637d627 ec1b552 4905514 ec1b552 db371b0 ec1b552 db371b0 ec1b552 8091bbd ec1b552 db371b0 ec1b552 8091bbd ec1b552 db371b0 ec1b552 637d627 92bf0a5 637d627 |
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 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 |
import gradio as gr
import subprocess
import os
import time
from transformers import AutoTokenizer, AutoModelForCausalLM
import logging
from starlette.middleware.sessions import SessionMiddleware
from fastapi import FastAPI
import uvicorn
app = FastAPI()
app.add_middleware(SessionMiddleware, secret_key="secure_key")
# Configure logging
logging.basicConfig(level=logging.INFO)
# Path to the cloned repository
BITNET_REPO_PATH = "/home/user/app/BitNet"
SETUP_SCRIPT = os.path.join(BITNET_REPO_PATH, "setup_env.py")
INFERENCE_SCRIPT = os.path.join(BITNET_REPO_PATH, "run_inference.py")
# Function to set up the environment by running setup.py
def setup_bitnet(model_name):
try:
result = subprocess.run(
f"python {SETUP_SCRIPT} --hf-repo {model_name} -q i2_s",
shell=True,
cwd=BITNET_REPO_PATH,
capture_output=True,
text=True
)
if result.returncode == 0:
return "Setup completed successfully!"
else:
return f"Error in setup: {result.stderr}"
except Exception as e:
return str(e)
# Function to run inference using the `run_inference.py` file
def run_inference(model_name, input_text, num_tokens=6):
try:
# Call the `run_inference.py` script with the model and input
start_time = time.time()
if input_text is None :
return "Please provide an input text for the model"
result = subprocess.run(
f"python run_inference.py -m models/{model_name}/ggml-model-i2_s.gguf -p \"{input_text}\" -n {num_tokens} -temp 0",
shell=True,
cwd=BITNET_REPO_PATH,
capture_output=True,
text=True
)
end_time = time.time()
if result.returncode == 0:
inference_time = round(end_time - start_time, 2)
return result.stdout, f"Inference took {inference_time} seconds."
else:
return f"Error during inference: {result.stderr}", None
except Exception as e:
return str(e), None
def run_transformers(profile: gr.OAuthProfile | None, oauth_token: gr.OAuthToken | None, model_name, input_text, num_tokens):
# if oauth_token is None :
# return "Error : To Compare please login to your HF account and make sure you have access to the used Llama models"
# Load the model and tokenizer dynamically if needed (commented out for performance)
tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=oauth_token.token)
model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=oauth_token.token)
if input_text is None :
return "Please provide an input text for the model", None
# Encode the input text
input_ids = tokenizer.encode(input_text, return_tensors="pt")
# Start time for inference
start_time = time.time()
# Generate output with the specified number of tokens
output = model.generate(input_ids, max_length=len(input_ids[0]) + num_tokens, num_return_sequences=1)
# Calculate inference time
inference_time = time.time() - start_time
# Decode the generated output
generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
return generated_text, f"{inference_time:.2f} seconds"
# Gradio Interface
def interface():
with gr.Blocks(css=".gr-button {background-color: #5C6BC0; color: white;} .gr-button:hover {background-color: #3F51B5;}") as demo:
# gr.LoginButton(elem_id="login-button", elem_classes="center-button")
gr.Markdown(
"""
<h1 style="text-align: center; color: #4A148C;">BitNet.cpp Speed Demonstration</h1>
<p style="text-align: center; color: #6A1B9A;">Compare the speed and performance of BitNet with Transformers!</p>
""",
elem_id="header"
)
# Model selection and setup row
with gr.Row():
model_dropdown = gr.Dropdown(
label="Select Model",
choices=["HF1BitLLM/Llama3-8B-1.58-100B-tokens", "1bitLLM/bitnet_b1_58-3B", "1bitLLM/bitnet_b1_58-large"], # Replace with available models
value="HF1BitLLM/Llama3-8B-1.58-100B-tokens",
interactive=True,
elem_id="model-dropdown"
)
setup_button = gr.Button("Run Setup", elem_id="setup-button")
setup_status = gr.Textbox(label="Setup Status", interactive=False, placeholder="Setup status will appear here...")
# Inference row
with gr.Row():
num_tokens = gr.Slider(minimum=1, maximum=100, label="Number of Tokens to Generate", value=50, step=1)
input_text = gr.Textbox(label="Input Text", placeholder="Enter your input text here...")
infer_button = gr.Button("Run Inference", elem_id="infer-button")
result_output = gr.Textbox(label="Output", interactive=False, placeholder="Inference output will appear here...")
time_output = gr.Textbox(label="Inference Time", interactive=False, placeholder="Inference time will appear here...")
# Comparison with Transformers
with gr.Row():
transformer_model_dropdown = gr.Dropdown(
label="Select Transformers Model",
choices=["TinyLlama/TinyLlama-1.1B-Chat-v1.0"], # Replace with actual models
value="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
interactive=True
)
compare_button = gr.Button("Run Transformers Inference", elem_id="compare-button")
transformer_result_output = gr.Textbox(label="Transformers Output", interactive=False, placeholder="Transformers output will appear here...")
transformer_time_output = gr.Textbox(label="Transformers Inference Time", interactive=False, placeholder="Transformers inference time will appear here...")
# Actions
setup_button.click(setup_bitnet, inputs=model_dropdown, outputs=setup_status)
infer_button.click(run_inference, inputs=[model_dropdown, input_text, num_tokens], outputs=[result_output, time_output])
compare_button.click(run_transformers, inputs=[transformer_model_dropdown, input_text, num_tokens], outputs=[transformer_result_output, transformer_time_output])
# Launch the Gradio app
return demo
demo = interface()
app.mount("/", demo)
# Launch the app
# demo.launch()
uvicorn.run(app, host="0.0.0.0", port=7860) |