from __future__ import annotations import os from dataclasses import dataclass from pathlib import Path import spaces import torch from transformers import AutoModelForCausalLM, AutoTokenizer from src.ppt.deck_generator import parse_slides_payload SYSTEM_PROMPT = ( "You are a presentation planner. Return ONLY a strict JSON array, no markdown and no commentary. " "Each array element must be an object with keys: " "'title' (string), 'bullets' (array of strings), and " "'layout_type' (one of: 'title_slide', 'content_slide', 'split_slide'). " "Use 5-8 slides, concise business-friendly bullets, and valid JSON syntax." ) @dataclass class GemmaClientConfig: model_id: str = "unsloth/gemma-4-12b-it-GGUF" transformers_model_id: str = "google/gemma-4-12b-it" temperature: float = 0.4 max_new_tokens: int = 900 backend: str = "transformers" # transformers | llama.cpp | vllm class GemmaClient: def __init__(self, config: GemmaClientConfig | None = None): self.config = config or GemmaClientConfig() self._tokenizer = None self._model = None def _load_transformers_model(self): if self._model is not None and self._tokenizer is not None: return # Patch transformers bug: Gemma-4 tokenizer sends extra_special_tokens as # a list but transformers calls .keys() on it expecting a dict. try: from transformers import tokenization_utils_base as _tub _orig_set = _tub.PreTrainedTokenizerBase._set_model_specific_special_tokens def _safe_set(self_tok, special_tokens): if isinstance(special_tokens, dict): _orig_set(self_tok, special_tokens) _tub.PreTrainedTokenizerBase._set_model_specific_special_tokens = _safe_set except Exception: pass model_id = self.config.transformers_model_id hf_token = os.getenv("HF_TOKEN") self._tokenizer = AutoTokenizer.from_pretrained( model_id, trust_remote_code=True, token=hf_token ) kwargs = { "device_map": "auto", "torch_dtype": torch.float16 if torch.cuda.is_available() else torch.float32, "trust_remote_code": True, "token": hf_token, } try: from transformers import BitsAndBytesConfig kwargs["quantization_config"] = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.float16, ) except Exception: # On unsupported environments (e.g., CPU/macOS without bitsandbytes), # fallback to non-quantized loading. pass self._model = AutoModelForCausalLM.from_pretrained(model_id, **kwargs) def _generate_transformers(self, user_prompt: str) -> str: self._load_transformers_model() messages = [ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": user_prompt}, ] raw = self._tokenizer.apply_chat_template( messages, tokenize=True, add_generation_prompt=True, return_tensors="pt", ) # Newer transformers may return a BatchEncoding; older returns a plain tensor if hasattr(raw, "input_ids"): input_ids = raw.input_ids.to(self._model.device) gen_kwargs = {"input_ids": input_ids} if hasattr(raw, "attention_mask"): gen_kwargs["attention_mask"] = raw.attention_mask.to(self._model.device) else: input_ids = raw.to(self._model.device) gen_kwargs = {"input_ids": input_ids} with torch.no_grad(): outputs = self._model.generate( **gen_kwargs, max_new_tokens=self.config.max_new_tokens, temperature=self.config.temperature, do_sample=True, pad_token_id=self._tokenizer.eos_token_id, ) prompt_len = input_ids.shape[1] generated_ids = outputs[0][prompt_len:] return self._tokenizer.decode(generated_ids, skip_special_tokens=True).strip() def _generate_llama_cpp(self, user_prompt: str) -> str: try: from llama_cpp import Llama except Exception as exc: raise RuntimeError( f"llama.cpp backend import failed: {exc}" ) from exc model_path = self._resolve_llama_cpp_model_path() llm = Llama( model_path=model_path, n_ctx=4096, n_gpu_layers=-1, verbose=False, ) completion = llm.create_chat_completion( messages=[ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": user_prompt}, ], temperature=self.config.temperature, max_tokens=self.config.max_new_tokens, ) return completion["choices"][0]["message"]["content"].strip() def _resolve_llama_cpp_model_path(self) -> str: env_model_path = os.getenv("LLAMA_CPP_MODEL_PATH") if env_model_path: return env_model_path preferred = [ "Q4_K_M.gguf", "q4_k_m.gguf", "Q4_K_S.gguf", "q4_k_s.gguf", "Q5_K_M.gguf", "q5_k_m.gguf", ] try: from huggingface_hub import HfApi, hf_hub_download except Exception as exc: raise RuntimeError( "llama.cpp backend requires either LLAMA_CPP_MODEL_PATH or huggingface_hub to auto-download GGUF." ) from exc repo_id = os.getenv("GGUF_MODEL_REPO", self.config.model_id) files = HfApi(token=os.getenv("HF_TOKEN")).list_repo_files(repo_id=repo_id) gguf_files = [f for f in files if f.lower().endswith(".gguf")] if not gguf_files: raise RuntimeError(f"No GGUF files found in repo: {repo_id}") selected = None for suffix in preferred: selected = next((f for f in gguf_files if f.endswith(suffix)), None) if selected: break if selected is None: selected = gguf_files[0] local_dir = Path(os.getenv("GGUF_CACHE_DIR", "models")) local_dir.mkdir(parents=True, exist_ok=True) downloaded = hf_hub_download( repo_id=repo_id, filename=selected, local_dir=str(local_dir), token=os.getenv("HF_TOKEN"), ) os.environ["LLAMA_CPP_MODEL_PATH"] = downloaded return downloaded def _generate_vllm(self, user_prompt: str) -> str: try: from openai import OpenAI except Exception as exc: raise RuntimeError("vLLM backend requested but openai package is unavailable.") from exc base_url = os.getenv("VLLM_BASE_URL", "http://localhost:8000/v1") model_name = os.getenv("VLLM_MODEL_NAME", self.config.model_id) client = OpenAI(base_url=base_url, api_key=os.getenv("VLLM_API_KEY", "EMPTY")) response = client.chat.completions.create( model=model_name, messages=[ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": user_prompt}, ], temperature=self.config.temperature, max_tokens=self.config.max_new_tokens, ) return response.choices[0].message.content.strip() @spaces.GPU def generate_json(self, user_prompt: str) -> str: backend = self.config.backend.lower().strip() if backend == "llama.cpp": return self._generate_llama_cpp(user_prompt) if backend == "vllm": return self._generate_vllm(user_prompt) return self._generate_transformers(user_prompt) def generate_slides(self, user_prompt: str) -> list[dict]: raw_output = self.generate_json(user_prompt) return parse_slides_payload(raw_output) _DEFAULT_CLIENT: GemmaClient | None = None def load_model(backend: str | None = None) -> GemmaClient: global _DEFAULT_CLIENT if _DEFAULT_CLIENT is None or (backend and _DEFAULT_CLIENT.config.backend != backend): cfg = GemmaClientConfig(backend=backend or os.getenv("MODEL_BACKEND", "transformers")) _DEFAULT_CLIENT = GemmaClient(cfg) return _DEFAULT_CLIENT def should_startup_warmup() -> bool: policy = os.getenv("ENABLE_STARTUP_WARMUP", "auto").strip().lower() if policy in {"1", "true", "yes", "on"}: return True if policy in {"0", "false", "no", "off"}: return False # Auto mode: enable warmup on hosted Space environments. return bool( os.getenv("SPACE_ID") or os.getenv("HF_SPACE_ID") or os.getenv("SPACE_AUTHOR_NAME") ) def warmup_model_cache(backend: str | None = None) -> str: client = load_model(backend=backend) chosen_backend = (backend or client.config.backend).lower().strip() if chosen_backend == "llama.cpp": model_path = client._resolve_llama_cpp_model_path() return f"llama.cpp cache ready: {model_path}" if chosen_backend == "transformers": client._load_transformers_model() return f"transformers model ready: {client.config.transformers_model_id}" if chosen_backend == "vllm": return "vllm warmup skipped: remote endpoint assumed" return f"warmup skipped: unsupported backend '{chosen_backend}'" def generate_slides(user_input: str, backend: str | None = None): client = load_model(backend=backend) return client.generate_slides(user_input) def generate_slides_json(user_input: str, backend: str | None = None) -> str: client = load_model(backend=backend) return client.generate_json(user_input) def safe_parse_generated_json(raw_output: str) -> list[dict]: return parse_slides_payload(raw_output)