| import os |
| os.environ["TOKENIZERS_PARALLELISM"] = "false" |
|
|
| import spaces |
| import json |
| from threading import Thread, Lock |
| from typing import Optional |
| from quota_guard import QuotaBlocked, verify_free_quota |
| from gguf_runtime import generate_gguf |
| import gradio as gr |
| import torch |
| from huggingface_hub import hf_hub_download |
| from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer |
|
|
| MODEL_ID = "SEE7/Astheral" |
| MODEL_REVISION = "37adef9b3917e4a3de0bb77f0db71089b7b82202" |
| settings = json.load(open(hf_hub_download(MODEL_ID, "settings.json", revision=MODEL_REVISION))) |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, revision=MODEL_REVISION) |
| model = AutoModelForCausalLM.from_pretrained( |
| MODEL_ID, revision=MODEL_REVISION, torch_dtype=torch.bfloat16, |
| device_map={"": "cuda"}, low_cpu_mem_usage=True, |
| attn_implementation="sdpa", use_safetensors=True, |
| ).eval() |
|
|
| @spaces.GPU(duration=120) |
| def _generate(message, history, max_tokens, temperature, mode, request: gr.Request): |
| if not message or not message.strip(): |
| raise gr.Error("Enter a message first.") |
| if len(message) > 32000: |
| raise gr.Error("Please keep your message under 32,000 characters.") |
| system = next((item["content"] for item in reversed(history) if item.get("role") == "system" and isinstance(item.get("content"), str)), settings["system"]) |
| if mode == "Thinking": |
| system += "\n" + settings["thinking"] |
| messages = [{"role": "system", "content": system}] |
| |
| turns = [item for item in history[-12:] if item.get("role") in ("user", "assistant") and isinstance(item.get("content"), str)] |
| latest = {"role": "user", "content": message.strip()} |
| while turns and len(tokenizer.apply_chat_template(messages + turns + [latest], tokenize=True, add_generation_prompt=True)) > 8192: |
| turns = turns[2:] |
| inputs = tokenizer.apply_chat_template( |
| messages + turns + [latest], tokenize=True, add_generation_prompt=True, |
| return_dict=True, return_tensors="pt", |
| ) |
| if inputs["input_ids"].shape[1] > 8192: |
| raise gr.Error("That message is too long for this budget-friendly demo. Please shorten it.") |
| inputs = inputs.to("cuda") |
| streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=115) |
| generation = dict( |
| **inputs, streamer=streamer, max_new_tokens=min(max(int(max_tokens), 1), 4096), |
| do_sample=float(temperature) > 0, max_time=110, |
| pad_token_id=128004, eos_token_id=128009, |
| ) |
| if float(temperature) > 0: |
| generation.update(temperature=float(temperature), top_p=0.9) |
| errors = [] |
| def generate(): |
| try: |
| with torch.inference_mode(): |
| model.generate(**generation) |
| except Exception as exc: |
| errors.append(exc) |
| streamer.end() |
| worker = Thread(target=generate, daemon=True) |
| worker.start() |
| response = "" |
| try: |
| for text in streamer: |
| response += text |
| yield response |
| finally: |
| worker.join(timeout=2) |
| if errors: |
| raise gr.Error("Generation failed. Please retry with a shorter prompt.") from errors[0] |
| if not response.strip(): |
| raise gr.Error("No text was generated. Please try again.") |
|
|
| quota_lock = Lock() |
|
|
| def respond(message, history, max_tokens, temperature, mode, request: gr.Request): |
| |
| with quota_lock: |
| try: |
| verify_free_quota() |
| except QuotaBlocked as exc: |
| raise gr.Error(str(exc)) from None |
| yield from _generate(message, history, max_tokens, temperature, mode, request) |
|
|
| def respond_gguf(model_id, message, history, max_tokens, temperature, request: gr.Request): |
| from gguf_models import model_path |
| from gguf_runtime import generate_gguf |
| with quota_lock: |
| try: |
| verify_free_quota() |
| path = model_path(model_id) |
| except (QuotaBlocked, ValueError) as exc: |
| raise gr.Error(str(exc)) from None |
| if not isinstance(message, str) or not message.strip() or len(message) > 32000: |
| raise gr.Error("Enter a message under 32,000 characters.") |
| messages = [item for item in history[-12:] if item.get("role") in ("system", "user", "assistant") and isinstance(item.get("content"), str)] |
| messages.append({"role": "user", "content": message}) |
| try: |
| yield from generate_gguf(path, messages, min(max(int(max_tokens), 1), 4096), float(temperature)) |
| except Exception as exc: |
| raise gr.Error("GGUF generation failed: " + str(exc)[:300]) from None |
|
|
| with gr.Blocks(theme=gr.themes.Soft(primary_hue="indigo"), title="Astheral · Built with Llama") as demo: |
| with gr.Accordion("Response settings", open=False): |
| max_tokens = gr.Slider(32, 4096, value=1024, step=32, label="Maximum new tokens") |
| temperature = gr.Slider(0, 1, value=0.6, step=0.1, label="Temperature (0 = deterministic)") |
| mode = gr.Radio(["Instruct", "Thinking"], value="Instruct", label="Mode") |
| gr.ChatInterface( |
| fn=respond, type="messages", additional_inputs=[max_tokens, temperature, mode], |
| examples=[["Explain Python generators in three sentences.", 192, 0.6, "Instruct"], ["Write a short starship announcement.", 192, 0.6, "Instruct"]], |
| cache_examples=False, concurrency_limit=1, api_name="chat", |
| ) |
| with demo: |
| gguf_id = gr.Textbox(visible=False) |
| gguf_message = gr.Textbox(visible=False) |
| gguf_history = gr.JSON(visible=False) |
| gguf_tokens = gr.Number(value=1024, visible=False) |
| gguf_temperature = gr.Number(value=0.6, visible=False) |
| gguf_output = gr.Textbox(visible=False) |
| gguf_trigger = gr.Button(visible=False) |
| gguf_trigger.click(respond_gguf, [gguf_id, gguf_message, gguf_history, gguf_tokens, gguf_temperature], |
| gguf_output, api_name="gguf_chat", concurrency_limit=1) |
| demo.queue(max_size=6, default_concurrency_limit=1) |
| if __name__ == "__main__": |
| from llama_ui import install_ui |
| demo.launch(prevent_thread_lock=True, show_error=True, ssr_mode=False) |
| install_ui(demo.app) |
| demo.block_thread() |
|
|
