Upload pipeline.py with huggingface_hub

pipeline.py

@@ -1,13 +1,11 @@
-import json
-import numpy as np
-from sentence_transformers import SentenceTransformer
+
+import json, torch, numpy as np
+from sentence_transformers import SentenceTransformer, CrossEncoder
 import faiss
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
-import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 class JujutsuKaiserver:
-    def __init__(self, model_dir="./upload_model"):
-        # Load config
+    def __init__(self, model_dir="."):
         with open(f"{model_dir}/rag_config.json") as f:
             config = json.load(f)
         self.embedder = SentenceTransformer(config["embedder_model"])
@@ -15,30 +13,33 @@ class JujutsuKaiserver:
         with open(f"{model_dir}/chunks.txt", "r", encoding="utf-8") as f:
             raw = f.read().split("<|CHUNK_END|>")
         self.chunks = [c.strip() for c in raw if c.strip()]
-        # Load model
-        bnb_config = BitsAndBytesConfig(
-            load_in_4bit=True,
-            bnb_4bit_use_double_quant=True,
-            bnb_4bit_quant_type='nf4',
-            bnb_4bit_compute_dtype=torch.float16
-        )
+        self.reranker = CrossEncoder(f"{model_dir}/cross_encoder_model")
         self.tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
         self.model = AutoModelForCausalLM.from_pretrained(
             model_dir,
-            quantization_config=bnb_config,
+            torch_dtype=torch.float16,
             device_map='auto',
             trust_remote_code=True
         )
+
     def ask(self, question, max_tokens=300):
+        q_lower = question.strip().lower()
+        if q_lower in ('hi', 'hello', 'hey', 'yo', 'sup', 'hi there'):
+            return "Hey there! I'm JujutsuKaiserver, your all-knowing JJK assistant. Ask me anything!"
         q_emb = self.embedder.encode([question]).astype('float32')
-        _, indices = self.index.search(q_emb, 5)
-        context = "\n\n".join([self.chunks[i] for i in indices[0]])
+        _, indices = self.index.search(q_emb, 30)
+        candidates = [self.chunks[i] for i in indices[0]]
+        pairs = [(question, c) for c in candidates]
+        scores = self.reranker.predict(pairs)
+        reranked = sorted(zip(scores, candidates), reverse=True)[:4]
+        best = [c for _, c in reranked]
+        context = "\n\n".join(best)
         messages = [
-            {"role": "system", "content": "You are JujutsuKaiserver, an expert on Jujutsu Kaisen. Answer using ONLY the context. If unsure, say you don't know."},
+            {"role": "system", "content": "You are JujutsuKaiserver, an expert on Jujutsu Kaisen. Answer using ONLY the provided context. Be friendly and concise."},
             {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {question}"}
         ]
         prompt = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
         inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
         outputs = self.model.generate(**inputs, max_new_tokens=max_tokens, temperature=0.7, do_sample=True, pad_token_id=self.tokenizer.eos_token_id)
         answer = self.tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
-        return answer.strip()
+        return answer.strip()