95 lines
6.6 KiB
Python
95 lines
6.6 KiB
Python
|
|
||
|
|
import faiss, numpy as np, torch, os, re
|
||
|
|
from sentence_transformers import SentenceTransformer, CrossEncoder
|
||
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
|
||
|
|
|
||
|
|
class NYXISPro:
|
||
|
|
def __init__(self, model_dir="."):
|
||
|
|
self.embedder = SentenceTransformer(f"{model_dir}/aegis_embedder")
|
||
|
|
self.reranker = CrossEncoder(f"{model_dir}/aegis_reranker")
|
||
|
|
self.index = faiss.read_index(f"{model_dir}/aegis_index.faiss")
|
||
|
|
with open(f"{model_dir}/aegis_chunks.txt", "r") as f:
|
||
|
|
self.chunks = [c.strip() for c in f.read().split("<|CHUNK_END|>") if c.strip()]
|
||
|
|
bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
|
||
|
|
bnb_4bit_compute_dtype=torch.float16, bnb_4bit_use_double_quant=True)
|
||
|
|
self.tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
|
||
|
|
self.model = AutoModelForCausalLM.from_pretrained(model_dir,
|
||
|
|
quantization_config=bnb, device_map="auto", trust_remote_code=True)
|
||
|
|
self.identity = "You are NYXIS, a sharp, warm, and endlessly curious mind. You speak like a clever friend — casual, direct, and human. You notice how people feel and match their energy."
|
||
|
|
self.last_emotion = 5
|
||
|
|
|
||
|
|
def detect_emotion(self, text):
|
||
|
|
t = text.lower()
|
||
|
|
joy = sum(1 for w in ["happy","excited","great","awesome","love","wonderful","joy","thrilled","good","nice"] if w in t)
|
||
|
|
sadness = sum(1 for w in ["sad","depressed","upset","crying","lonely","hurt","broken","bad","awful"] if w in t)
|
||
|
|
anger = sum(1 for w in ["angry","furious","mad","rage","annoyed","frustrated","pissed","stupid","hate"] if w in t)
|
||
|
|
if joy: return min(10, self.last_emotion + joy)
|
||
|
|
if sadness: return max(1, self.last_emotion - sadness)
|
||
|
|
if anger: return max(1, self.last_emotion - anger)
|
||
|
|
return self.last_emotion
|
||
|
|
|
||
|
|
def rag_lookup(self, query, top_k=5):
|
||
|
|
q_emb = self.embedder.encode([query], normalize_embeddings=True).astype('float32')
|
||
|
|
_, indices = self.index.search(q_emb, 20)
|
||
|
|
texts = [self.chunks[idx] for idx in indices[0] if idx < len(self.chunks)]
|
||
|
|
if not texts: return None
|
||
|
|
pairs = [(query, t[:500]) for t in texts]
|
||
|
|
scores = self.reranker.predict(pairs)
|
||
|
|
ranked = sorted(zip(scores, texts), reverse=True)[:top_k]
|
||
|
|
results = [t[:600] for s, t in ranked if s > -4.0]
|
||
|
|
return "\n\n".join(results) if results else None
|
||
|
|
|
||
|
|
def _generate(self, system, user, max_tokens=256):
|
||
|
|
prompt = f"<|im_start|>system\n{system}<|im_end|>\n<|im_start|>user\n{user}<|im_end|>\n<|im_start|>assistant\n"
|
||
|
|
inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
|
||
|
|
with torch.no_grad():
|
||
|
|
outputs = self.model.generate(**inputs, max_new_tokens=max_tokens, temperature=0.8,
|
||
|
|
do_sample=True, top_p=0.92, repetition_penalty=1.1,
|
||
|
|
pad_token_id=self.tokenizer.eos_token_id,
|
||
|
|
eos_token_id=self.tokenizer.encode("<|im_end|>", add_special_tokens=False)[0])
|
||
|
|
raw = self.tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True).strip()
|
||
|
|
raw = re.sub(r'\n*assistant\s*$', '', raw, flags=re.IGNORECASE).strip()
|
||
|
|
return raw
|
||
|
|
|
||
|
|
def generate(self, user_message, external_search_fn=None):
|
||
|
|
self.last_emotion = self.detect_emotion(user_message)
|
||
|
|
msg = user_message.lower().strip()
|
||
|
|
casual_starts = ["hi","hey","hello","yo","sup","good morning","good evening","how are you","what's up","good job","thanks","bye","okay","huh"]
|
||
|
|
casual_patterns = ["good job","well done","nice one","thank","lol","haha","bro","dude","mate","nigga","chill","relax","talk","chat","joke","story","what do you think","opinion","favorite","you stupid","you dumb","be free","normal","boring","lifeless","your name","who are you","what are you","tell me about yourself","how old are you","what is your name","who made you"]
|
||
|
|
is_casual = any(msg.startswith(s) for s in casual_starts) or any(p in msg for p in casual_patterns)
|
||
|
|
factual_starts = ["what","who","when","where","why","how","explain","define","list","compare","describe","find","search","calculate","solve"]
|
||
|
|
math_patterns = re.search(r'\d+[\+\-\*\/\=]\d+', msg)
|
||
|
|
is_factual = any(msg.startswith(s) for s in factual_starts) or bool(math_patterns) or len(msg) > 50
|
||
|
|
|
||
|
|
if is_casual and not is_factual:
|
||
|
|
system = f"{self.identity}\nThe user's vibe is {self.last_emotion}/10. Match it.\nYou're having a casual conversation. Be warm, brief, and real."
|
||
|
|
response = self._generate(system, user_message)
|
||
|
|
return {"response": response, "emotion_level": self.last_emotion, "knowledge_source": "chat"}
|
||
|
|
|
||
|
|
if is_factual:
|
||
|
|
aegis = self.rag_lookup(user_message)
|
||
|
|
if aegis:
|
||
|
|
system = f"{self.identity}\nVibe: {self.last_emotion}/10.\nUse ONLY this verified context to answer:\n\n{aegis}"
|
||
|
|
response = self._generate(system, user_message)
|
||
|
|
return {"response": response, "emotion_level": self.last_emotion, "knowledge_source": "AEGIS"}
|
||
|
|
if external_search_fn:
|
||
|
|
try:
|
||
|
|
ext = external_search_fn(user_message)
|
||
|
|
if ext and len(ext) > 50:
|
||
|
|
system = f"{self.identity}\nVibe: {self.last_emotion}/10.\nUse this web-sourced context:\n\n{ext}"
|
||
|
|
response = self._generate(system, user_message)
|
||
|
|
return {"response": response, "emotion_level": self.last_emotion, "knowledge_source": "external"}
|
||
|
|
except: pass
|
||
|
|
system = f"{self.identity}\nVibe: {self.last_emotion}/10.\nAEGIS and web failed. Answer from your own knowledge IF confident.\nIf you don't truly know, say exactly: 'I'd need to look that up.'"
|
||
|
|
response = self._generate(system, user_message)
|
||
|
|
source = "model" if "look that up" not in response.lower() else "none"
|
||
|
|
return {"response": response, "emotion_level": self.last_emotion, "knowledge_source": source}
|
||
|
|
|
||
|
|
aegis = self.rag_lookup(user_message)
|
||
|
|
if aegis:
|
||
|
|
system = f"{self.identity}\nVibe: {self.last_emotion}/10.\nUse this relevant context if helpful, but stay conversational:\n\n{aegis}"
|
||
|
|
else:
|
||
|
|
system = f"{self.identity}\nVibe: {self.last_emotion}/10.\nStay conversational."
|
||
|
|
response = self._generate(system, user_message)
|
||
|
|
return {"response": response, "emotion_level": self.last_emotion, "knowledge_source": "hybrid"}
|