Upload pipeline.py with huggingface_hub

pipeline.py

@@ -0,0 +1,94 @@
+
+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"}