178 lines
11 KiB
Python
178 lines
11 KiB
Python
"""Pinned offline EmbeddingGemma 2 GGUF worker. No downloads or HTTP inference."""
|
|
import atexit
|
|
import base64
|
|
import hashlib
|
|
import io
|
|
import json
|
|
import math
|
|
import os
|
|
from pathlib import Path
|
|
import queue
|
|
import re
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
import wave
|
|
|
|
MODEL_ID="google/embeddinggemma-2"
|
|
REVISION="bfcd298762cc34d0357ece5ebdd31791a3a374d8"
|
|
DIMENSIONS=768
|
|
PROFILE="document-v1"
|
|
RUNTIME_REVISION="b11461"
|
|
RUNTIME_BUILD="b11461-no-logits-v1"
|
|
MAX_CHUNKS=64
|
|
MAX_LINE=4*1024*1024
|
|
MAX_REQUEST=16*1024*1024
|
|
MODEL_FILES={
|
|
"embeddinggemma-2-Q8_0.gguf":(309855456,"2188ac1deca4b77dffefd603c2776a9d76d9d74ec01841392982ebb840b09135"),
|
|
"mmproj-embeddinggemma-2-Q8_0.gguf":(554821024,"c4a8a52691ecef40618438928bdf9e68379b854e24166f292592353db0aab64f")}
|
|
|
|
|
|
def verify_file(directory,name,size,digest):
|
|
if not isinstance(name,str) or not name or Path(name).is_absolute() or ".." in Path(name).parts: raise ValueError("Invalid local model path")
|
|
file=(directory/name).resolve()
|
|
if not file.is_relative_to(directory) or not file.is_file() or file.stat().st_size!=size: raise ValueError("Local semantic artifact unavailable")
|
|
if not re.fullmatch(r"[0-9a-f]{64}",digest): raise ValueError("Invalid artifact digest")
|
|
calculated=hashlib.sha256()
|
|
with file.open("rb") as source:
|
|
while part:=source.read(1024*1024): calculated.update(part)
|
|
if calculated.hexdigest()!=digest: raise ValueError("Local semantic artifact verification failed")
|
|
return file
|
|
|
|
|
|
def validate_vectors(vectors,expected=None):
|
|
if not isinstance(vectors,list) or not 1<=len(vectors)<=MAX_CHUNKS or expected is not None and len(vectors)!=expected: raise ValueError("Invalid semantic vector batch")
|
|
for vector in vectors:
|
|
if not isinstance(vector,list) or len(vector)!=DIMENSIONS or any(not isinstance(value,(int,float)) or isinstance(value,bool) or not math.isfinite(value) for value in vector): raise ValueError("Invalid semantic output")
|
|
norm=math.sqrt(sum(value*value for value in vector))
|
|
if not 0.99<=norm<=1.01: raise ValueError("Semantic output is not normalized")
|
|
return vectors
|
|
|
|
|
|
def validate_readiness(ready):
|
|
if not isinstance(ready,dict) or ready.get("ready") is not True or ready.get("model")!=MODEL_ID or ready.get("dimensions")!=DIMENSIONS or ready.get("runtime")!=RUNTIME_BUILD or ready.get("device")!="cpu" or ready.get("vulkan_module_loaded") is not False:
|
|
raise ValueError("Local semantic worker identity mismatch")
|
|
|
|
|
|
class NativeSemanticModel:
|
|
def __init__(self,directory,worker_path):
|
|
self.directory=Path(directory).resolve(); self.worker_path=Path(worker_path).resolve()
|
|
manifest_path=self.directory/"model-manifest.json"
|
|
if not manifest_path.is_file() or manifest_path.stat().st_size>65536: raise ValueError("Local model manifest unavailable")
|
|
manifest=json.loads(manifest_path.read_text(encoding="utf-8"))
|
|
if manifest.get("model_id")!=MODEL_ID or manifest.get("revision")!=REVISION or manifest.get("dimensions")!=DIMENSIONS or manifest.get("quantization")!="Q8_0" or manifest.get("runtime_tag")!=RUNTIME_REVISION: raise ValueError("Pinned Google model identity mismatch")
|
|
self.model=verify_file(self.directory,*next((name,*value) for name,value in MODEL_FILES.items() if not name.startswith("mmproj")))
|
|
self.projector=verify_file(self.directory,*next((name,*value) for name,value in MODEL_FILES.items() if name.startswith("mmproj")))
|
|
runtime_path=self.worker_path.parent/"runtime-manifest.json"
|
|
if not runtime_path.is_file() or runtime_path.stat().st_size>65536: raise ValueError("Verified native runtime manifest unavailable")
|
|
runtime=json.loads(runtime_path.read_text(encoding="utf-8")); files=runtime.get("files")
|
|
if not isinstance(files,list) or not 1<=len(files)<=64 or not any(item.get("path")==self.worker_path.name for item in files): raise ValueError("Invalid native runtime manifest")
|
|
for item in files: verify_file(self.worker_path.parent,item["path"],item["bytes"],item["sha256"])
|
|
# All DLLs load from the private, verified runtime directory.
|
|
verified={item["path"].casefold() for item in files}
|
|
if any(file.name.casefold() not in verified for file in self.worker_path.parent.glob("*.dll")): raise ValueError("Unverified runtime module")
|
|
self.lock=threading.Lock();self.process=None;self.responses=None
|
|
atexit.register(self.close)
|
|
|
|
def info(self):
|
|
return {"model":MODEL_ID,"revision":REVISION,"dimensions":DIMENSIONS,"profile":PROFILE,"artifactFormat":"gguf","quantization":"Q8_0","runtimeRevision":RUNTIME_REVISION,"local":True,"backend":"native-gguf","device":"cpu","modalities":["text","image","audio"]}
|
|
|
|
def close(self):
|
|
process=self.process;self.process=None
|
|
if process is not None:
|
|
try: process.kill();process.wait(timeout=3)
|
|
except Exception: pass
|
|
for stream in (process.stdin,process.stdout):
|
|
try: stream.close()
|
|
except Exception: pass
|
|
|
|
def _read(self,process,responses):
|
|
try:
|
|
while True:
|
|
line=process.stdout.readline(MAX_LINE+1)
|
|
if not line or len(line)>MAX_LINE or not line.endswith(b"\n"): raise ValueError()
|
|
response=json.loads(line)
|
|
responses.put(response,timeout=1)
|
|
except Exception:
|
|
try: responses.put({"protocolFailure":True},timeout=1)
|
|
except Exception: pass
|
|
|
|
def _receive(self,timeout):
|
|
try:
|
|
response=self.responses.get(timeout=timeout)
|
|
if not isinstance(response,dict) or response.get("protocolFailure"): raise ValueError()
|
|
return response
|
|
except Exception:
|
|
self.close();raise ValueError("Local semantic worker failed or timed out") from None
|
|
|
|
def _start(self):
|
|
if self.process is not None and self.process.poll() is None: return
|
|
self.close();self.responses=queue.Queue(maxsize=2)
|
|
environment={key:value for key,value in os.environ.items() if key.upper() in {"SYSTEMROOT","WINDIR","TEMP","TMP"}}
|
|
environment.update({"SEARCHI_FORCE_CPU":"1","GGML_CPU_ONLY":"1","GGML_DISABLE_VULKAN":"1"})
|
|
self.process=subprocess.Popen([str(self.worker_path),"--model",str(self.model),"--threads","2","--device","cpu"],
|
|
cwd=str(self.worker_path.parent),stdin=subprocess.PIPE,stdout=subprocess.PIPE,stderr=subprocess.DEVNULL,
|
|
env=environment,creationflags=getattr(subprocess,"CREATE_NO_WINDOW",0))
|
|
threading.Thread(target=self._read,args=(self.process,self.responses),daemon=True).start()
|
|
ready=self._receive(20)
|
|
try: validate_readiness(ready)
|
|
except Exception: self.close();raise
|
|
|
|
def _request(self,request,timeout=30):
|
|
payload=json.dumps(request,ensure_ascii=False,allow_nan=False).encode("utf-8")+b"\n"
|
|
if len(payload)>=MAX_REQUEST: raise ValueError("Semantic request budget exceeded")
|
|
with self.lock:
|
|
self._start()
|
|
process=self.process
|
|
# A blocked pipe write is also bounded and tears down the process.
|
|
timer=threading.Timer(timeout,lambda:process.kill() if process.poll() is None else None);timer.daemon=True;timer.start()
|
|
try:
|
|
process.stdin.write(payload);process.stdin.flush();response=self._receive(timeout)
|
|
if response.get("ok") is not True or response.get("device")!="cpu": raise ValueError("Local semantic inference failed")
|
|
return response
|
|
except Exception:
|
|
self.close();raise ValueError("Local semantic inference failed") from None
|
|
finally: timer.cancel();payload=None
|
|
|
|
def text_vectors(self,text):
|
|
if not isinstance(text,str) or not text.strip() or len(text.encode("utf-8"))>1024*1024: raise ValueError("Semantic text budget exceeded")
|
|
encoded=text.encode("utf-8");windows=self._request({"op":"windows","text":text,"title":"none"}).get("windows")
|
|
if not isinstance(windows,list) or not 1<=len(windows)<=MAX_CHUNKS: raise ValueError("Semantic window budget exceeded")
|
|
previous_end=0;parts=[]
|
|
for item in windows:
|
|
start=item.get("start_byte");end=item.get("end_byte");part=item.get("text")
|
|
if not isinstance(start,int) or not isinstance(end,int) or not isinstance(part,str) or start<0 or start>previous_end or end<=previous_end or end>len(encoded) or part.encode("utf-8")!=encoded[start:end]: raise ValueError("Semantic source coverage invalid")
|
|
previous_end=end;parts.append(part)
|
|
if previous_end!=len(encoded): raise ValueError("Semantic source tail missing")
|
|
vectors=[];deadline=time.monotonic()+60
|
|
for index in range(0,len(parts),8):
|
|
remaining=deadline-time.monotonic()
|
|
if remaining<=0: raise ValueError("Semantic inference budget exceeded")
|
|
batch=parts[index:index+8];response=self._request({"op":"documents","texts":batch,"title":"none"},min(30,remaining))
|
|
vectors.extend(validate_vectors(response.get("vectors"),len(batch)))
|
|
return vectors
|
|
|
|
def image_vectors(self,data):
|
|
from PIL import Image,ImageSequence,ImageOps
|
|
image=Image.open(io.BytesIO(data));frames=getattr(image,"n_frames",1)
|
|
if frames>32 or image.width*image.height*frames>50_000_000: raise ValueError("Semantic image budget exceeded")
|
|
vectors=[]
|
|
for frame in ImageSequence.Iterator(image):
|
|
frame=ImageOps.exif_transpose(frame).convert("RGBA")
|
|
background=Image.new("RGBA",frame.size,"white");background.alpha_composite(frame);frame=background.convert("RGB")
|
|
scale=min(1,672/max(frame.size));width=max(48,round(frame.width*scale/48)*48);height=max(48,round(frame.height*scale/48)*48)
|
|
frame=frame.resize((width,height),Image.Resampling.LANCZOS)
|
|
request={"op":"image","projector":str(self.projector),"width":width,"height":height,"data":base64.b64encode(frame.tobytes()).decode("ascii")}
|
|
vectors.extend(validate_vectors(self._request(request).get("vectors"),1))
|
|
return vectors
|
|
|
|
def audio_vectors(self,data):
|
|
import numpy as np
|
|
with wave.open(io.BytesIO(data),"rb") as source:
|
|
if source.getnchannels()!=1 or source.getframerate()!=16000 or source.getsampwidth()!=2 or not 0<source.getnframes()<=480000: raise ValueError("Only bounded mono16kHz PCM16 WAV is supported")
|
|
payload=source.readframes(source.getnframes())
|
|
if len(payload)!=source.getnframes()*2: raise ValueError("Truncated audio data")
|
|
pcm=np.frombuffer(payload,dtype="<i2").astype("<f4")/32768
|
|
request={"op":"audio","projector":str(self.projector),"data":base64.b64encode(pcm.tobytes()).decode("ascii")}
|
|
return validate_vectors(self._request(request).get("vectors"),1)
|