Bonsai-27B / app-old.py
Awesome-Developer's picture
Rename app.py to app-old.py
a1f422a verified
Raw
History Blame Contribute Delete
2.38 kB
import os
import sys
import subprocess
# 1. FORCE EXTRACT PRE-BUILT CUDA WHEELS INSTANTLY BEFORE ANY IMPORTS
try:
import llama_cpp
except ModuleNotFoundError:
print("Package missing! Forcing immediate 5-second pre-compiled wheel installation...")
subprocess.check_call([
sys.executable, "-m", "pip", "install",
"llama-cpp-python",
"--extra-index-url", "https://github.io"
])
# Enforce cache directories to live on persistent storage
os.environ["HF_HOME"] = "/data/.huggingface"
os.environ["XDG_CACHE_HOME"] = "/data/.cache"
# CRITICAL FIX: Explicitly import hf_hub_download to fix the NameError
from huggingface_hub import hf_hub_download
from llama_cpp import Llama
import spaces
import gradio as gr
# 2. Pull down the standard 1-bit 3.9 GB file variant to your persistent bucket mount
print("Checking persistent storage for Bonsai-27B weights...")
model_path = hf_hub_download(
repo_id="prism-ml/Bonsai-27B-gguf",
filename="Bonsai-27B-Q1_0.gguf",
local_dir="/data"
)
# 3. Load the model into memory globally
print("Initializing model architecture into global memory...")
llm = Llama(
model_path=model_path,
n_ctx=32768, # Your 32K context window sweet spot
n_gpu_layers=99, # Dumps 100% of the layers directly to ZeroGPU VRAM
ctk_quant=2, # Compresses attention Keys to 4-bit format
ctv_quant=2, # Compresses attention Values to 4-bit format
flash_attn=True
)
# 4. ZeroGPU Endpoint execution loop
@spaces.GPU(duration=60)
def api_generate(prompt):
clean_prompt = str(prompt).strip()
if not clean_prompt:
return "Error: Empty prompt received."
formatted_prompt = f"<|user|>\n{clean_prompt}<|endoftext|>\n<|assistant|>"
response = llm(
formatted_prompt,
max_tokens=1024,
temperature=0.7,
top_p=0.95
)
try:
return response["choices"][0]["text"] # Fixed nested array dictionary index format
except (KeyError, IndexError, TypeError):
return str(response)
# 5. Bind the public endpoint mapping
demo = gr.Interface(
fn=api_generate,
inputs=gr.Textbox(label="Input Prompt / Document Dump", lines=8),
outputs=gr.Textbox(label="Model Output"),
api_name="chat" # Exposes your endpoint hook as /predict
)
if __name__ == "__main__":
demo.launch()