Jarvis / brain.py

Upload 11 files

6645e8f verified 6 days ago

5.49 kB

# 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

	# 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