Spaces:
Running on Zero
Running on Zero
| 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 | |
| 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() | |