File size: 4,663 Bytes
7c67638 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 | """
Thin HuggingFace-style wrapper around the webmind-brain engine.
Usage:
from webmind import Brain
brain = Brain.from_pretrained("webmind/webmind-brain-v1")
brain.ask("capital of france")
"""
import os
import sys
import json
import shutil
from pathlib import Path
def _find_engine():
"""Locate the actual brain engine source."""
# Check common locations
candidates = [
Path(__file__).parent.parent / "engine" / "src",
Path(__file__).parent.parent.parent / "webmind-research" / "papers" / "new-gen-ai" / "src",
Path.home() / "webmind-research" / "papers" / "new-gen-ai" / "src",
]
for p in candidates:
if (p / "brain.py").exists():
return str(p)
# Check WEBMIND_ENGINE_PATH env var
env_path = os.environ.get("WEBMIND_ENGINE_PATH")
if env_path and Path(env_path).exists():
return env_path
return None
class Brain:
"""HuggingFace-compatible wrapper for webmind-brain."""
def __init__(self, db_path=None):
engine_path = _find_engine()
if engine_path is None:
raise ImportError(
"Could not find webmind-brain engine source. "
"Set WEBMIND_ENGINE_PATH to the 'src/' directory, or install from the repo."
)
if engine_path not in sys.path:
sys.path.insert(0, engine_path)
from brain import Brain as BrainEngine
self._engine = BrainEngine(db_path=db_path)
@classmethod
def from_pretrained(cls, model_id: str, db_path: str = None, **kwargs):
"""
Load a brain, optionally downloading from HuggingFace Hub.
Args:
model_id: HuggingFace model ID (e.g. "webmind/webmind-brain-v1")
or local path to a brain database directory.
db_path: Override path for the neuron database.
"""
resolved_path = db_path
# If model_id is a local directory with a neurons.db, use it directly
if os.path.isdir(model_id) and os.path.exists(os.path.join(model_id, "neurons.db")):
resolved_path = resolved_path or model_id
elif db_path is None:
# Default: use ~/.cache/webmind/<model_name>
cache_dir = Path.home() / ".cache" / "webmind" / model_id.replace("/", "--")
cache_dir.mkdir(parents=True, exist_ok=True)
resolved_path = str(cache_dir)
# Try downloading from HF Hub if not cached
if not (cache_dir / "neurons.db").exists():
try:
from huggingface_hub import snapshot_download
downloaded = snapshot_download(model_id, cache_dir=str(cache_dir.parent))
# Copy db files to cache
for f in Path(downloaded).glob("*.db*"):
shutil.copy2(f, cache_dir / f.name)
except Exception:
# No download available — start with empty brain
pass
return cls(db_path=resolved_path)
def ask(self, question: str) -> dict:
"""
Ask the brain a question.
Returns dict with keys:
answer: str — the answer text
confidence: float — how confident (0-1)
strategy: str — how it was answered (convergence/co-occurrence/abstain)
trace: str — reasoning trace
"""
return self._engine.ask(question)
def generate(self, prompt: str, max_tokens: int = 30, temperature: float = 0.7) -> dict:
"""
Generate fluent text steered by a prompt.
Returns dict with keys:
text: str — generated text
trace: list — per-token generation trace
tokens_generated: int
"""
return self._engine.generate(prompt, max_tokens=max_tokens, temperature=temperature)
def teach(self, sentence: str, confidence: float = 0.6) -> list:
"""Teach the brain a new fact. Returns neuron IDs."""
return self._engine.teach(sentence, confidence=confidence)
def teach_batch(self, sentences: list, confidence: float = 0.6) -> list:
"""Teach multiple sentences at once."""
return self._engine.teach_batch(sentences, confidence=confidence)
def flush(self):
"""Persist all pending teaches to disk."""
self._engine.flush()
def health(self) -> dict:
"""Get brain health metrics."""
return self._engine.health()
def close(self):
"""Clean shutdown."""
self._engine.close()
def __enter__(self):
return self
def __exit__(self, *args):
self.close()
|