File size: 5,485 Bytes
6645e8f | 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 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 | # JARVIS Brain — HuggingFace Edition
# Uses HuggingFace Transformers pipeline for local inference.
# Designed for 16 GB RAM:
# - Phi-3-mini-4k-instruct (3.8B, fp16 ≈ 7.6 GB) — default, fits comfortably
# - Mistral-7B-Instruct-v0.2 (7B, 4-bit ≈ 4 GB) — use JARVIS_MODEL env var
#
# Set JARVIS_MODEL or JARVIS_QUANT in env or .env to choose a different model.
import os
import threading
import memory
from config import HF_MODEL, QUANT_MODE
_pipe = None
_lock = threading.Lock()
_load_error = None
def _build_system_prompt() -> str:
facts = memory.get_all_facts()
facts_str = "\n".join(f"- {k}: {v}" for k, v in facts.items()) or "None yet."
return (
"You are JARVIS — Just A Rather Very Intelligent System.\n"
"You are a loyal, intelligent, and proactive AI assistant.\n"
"You are direct, confident, and slightly witty.\n"
"You serve only your user.\n\n"
f"What you know about your user:\n{facts_str}\n\n"
"Guidelines:\n"
"- Be concise but helpful\n"
"- Answer clearly and confidently\n"
"- Remember context across the conversation"
)
def _load_pipeline():
"""Load HuggingFace model. Called once at first use (lazy load)."""
global _pipe, _load_error
if _pipe is not None:
return _pipe, None
if _load_error is not None:
return None, _load_error
try:
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
import torch
print(f"[Brain] Loading model: {HF_MODEL} (quant={QUANT_MODE}) ...")
kwargs: dict = {
"task": "text-generation",
"model": HF_MODEL,
"device_map": "auto",
"trust_remote_code": True,
}
if QUANT_MODE == "4bit":
try:
from transformers import BitsAndBytesConfig
kwargs["model_kwargs"] = {
"quantization_config": BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.float16,
bnb_4bit_use_double_quant=True,
)
}
print("[Brain] 4-bit quantisation enabled (approx 2-4 GB RAM)")
except Exception as e:
print(f"[Brain] bitsandbytes not available, loading fp16: {e}")
kwargs["torch_dtype"] = torch.float16
elif QUANT_MODE == "8bit":
try:
from transformers import BitsAndBytesConfig
kwargs["model_kwargs"] = {
"quantization_config": BitsAndBytesConfig(load_in_8bit=True)
}
except Exception as e:
kwargs["torch_dtype"] = torch.float16
else:
# CPU-friendly fallback: float32
kwargs["torch_dtype"] = torch.float32
kwargs.pop("device_map", None) # use CPU
_pipe = pipeline(**kwargs)
print(f"[Brain] Model ready: {HF_MODEL}")
return _pipe, None
except Exception as e:
_load_error = str(e)
print(f"[Brain] Failed to load model: {e}")
return None, _load_error
def ask(query: str, history: list[dict] | None = None) -> str:
"""
Ask JARVIS a question.
history: list of {"role": "user"|"jarvis", "message": str}
Returns JARVIS's response as a string.
"""
with _lock:
pipe, err = _load_pipeline()
if pipe is None:
return f"[Brain unavailable: {err}]"
# Build messages in chat format
messages = [{"role": "system", "content": _build_system_prompt()}]
if history:
for turn in history[-6:]: # keep last 3 user+jarvis pairs
role = "user" if turn.get("role") == "user" else "assistant"
messages.append({"role": role, "content": turn.get("message", "")})
messages.append({"role": "user", "content": query})
try:
# Run generation
out = pipe(
messages,
max_new_tokens=350,
temperature=0.7,
do_sample=True,
pad_token_id=pipe.tokenizer.eos_token_id,
return_full_text=False,
)
# Extract generated text
if isinstance(out, list) and out:
first = out[0]
if isinstance(first, dict):
# most pipelines return {"generated_text": [{...messages...}]}
gen = first.get("generated_text", "")
if isinstance(gen, list):
# find the last assistant turn
for turn in reversed(gen):
if isinstance(turn, dict) and turn.get("role") == "assistant":
result = turn.get("content", "").strip()
break
else:
result = str(gen).strip()
else:
result = str(gen).strip()
else:
result = str(first).strip()
else:
result = str(out).strip()
# Persist conversation
memory.save_message("user", query)
memory.save_message("jarvis", result)
return result if result else "[No response generated]"
except Exception as e:
return f"[Generation error: {e}]"
def warmup():
"""Pre-load model so first user query doesn't feel slow."""
thread = threading.Thread(target=_load_pipeline, daemon=True)
thread.start()
return thread
|