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"}