import gradio as gr
import subprocess
import os
import time
from transformers import AutoTokenizer, AutoModelForCausalLM
import logging
from starlette.middleware.sessions import SessionMiddleware
# 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
model_name = model_name.split("/")[1]
start_time = time.time()
if input_text is None or input_text == "":
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(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)
# if model_name=="TinyLlama/TinyLlama-1.1B-Chat-v1.0" :
print(input_text)
if input_text is None or input_text == "":
return "Please provide an input text for the model", None
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
# 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(theme=gr.themes.Ocean()) as demo:
# Header
gr.Markdown(
"""
<h1 style="text-align: center; color: #7AB8E5;">BitNet.cpp Speed Demonstration 💻</h1>
<p style="text-align: center; color: #6A1B9A;">Compare the speed and performance of BitNet with popular Transformer models.</p>
""",
elem_id="header"
)
# Instructions
gr.Markdown(
"""
### Instructions for Using the BitNet.cpp Speed Demonstration
1. **Set Up Your Project**: Begin by selecting the model you wish to use. Please note that this process may take a few minutes to complete.
2. **Select Token Count**: Choose the number of tokens you want to generate for your inference.
3. **Input Your Text**: Enter the text you wish to analyze, then compare the performance of BitNet with popular Transformer models.
""",
elem_id="instructions"
)
# Model Selection and Setup
with gr.Column(elem_id="container"):
gr.Markdown("<h2 style='color: #5CA2D3; text-align: center;'>Model Selection and Setup</h2>")
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"
],
value="HF1BitLLM/Llama3-8B-1.58-100B-tokens",
interactive=True
)
setup_button = gr.Button("Run Setup")
setup_status = gr.Textbox(label="Setup Status", interactive=False, placeholder="Setup status will appear here...")
# Inference Section
with gr.Column(elem_id="container"):
gr.Markdown("<h2 style='color: #5CA2D3; text-align: center;'>BitNet Inference</h2>")
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...",
value="Who is Zeus?"
)
with gr.Row():
infer_button = gr.Button("Run Inference")
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 Section
with gr.Column(elem_id="container"):
gr.Markdown("<h2 style='color: #5CA2D3; text-align: center;'>Compare with Transformers</h2>")
with gr.Row():
transformer_model_dropdown = gr.Dropdown(
label="Select Transformers Model",
choices=["TinyLlama/TinyLlama_v1.1"],
value="TinyLlama/TinyLlama_v1.1",
interactive=True
)
input_text_tr = gr.Textbox(label="Input Text", placeholder="Enter your input text here...", value="Who is Zeus?")
with gr.Row():
compare_button = gr.Button("Run Transformers Inference")
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_tr, num_tokens], outputs=[transformer_result_output, transformer_time_output])
return demo
demo = interface()
# # Access FastAPI app instance from Gradio
# fastapi_app = demo.app
# # Add SessionMiddleware to enable session management
# fastapi_app.add_middleware(SessionMiddleware, secret_key="secret_key") # Use a secure, random secret key
# # Launch the app
demo.launch()
# from fastapi import FastAPI
# app = FastAPI()
# # Add SessionMiddleware for sessions handling
# app.add_middleware(SessionMiddleware, secret_key="secure_secret_key")
# # Mount Gradio app to FastAPI at the root
# app.mount("/", demo)