Bigger model

app.py

@@ -20,17 +20,18 @@ logger = logging.getLogger(__name__)
 
 app = FastAPI()
 
-# Authenticate with Hugging Fac
+# Authenticate with Hugging Face
 hf_token = os.getenv("HF_TOKEN")
 if not hf_token:
     logger.error("HF_TOKEN environment variable not set.")
     raise ValueError("HF_TOKEN not set")
 login(token=hf_token)
 
-# Models
+# Models Configuration
 sentence_transformer_model = "all-MiniLM-L6-v2"
-repo_id = "bartowski/deepcogito_cogito-v1-preview-llama-3B-GGUF"
-filename = "deepcogito_cogito-v1-preview-llama-3B-Q4_K_M.gguf"  # Updated to Cogito Q4_K_M
+# Upgrade to the 8B model and choose Q4_0 quantization for a good balance of performance and resource usage.
+repo_id = "bartowski/deepcogito_cogito-v1-preview-llama-8B-GGUF"
+filename = "deepcogito_cogito-v1-preview-llama-8B-Q4_KM.gguf"  # New 8B model with Q4_0 quantization
 
 # Define FAQs (unchanged)
 faqs = [
@@ -55,17 +56,17 @@ try:
     faiss_index.add(cv_embeddings)
     logger.info("FAISS index built successfully")
 
-    # Load embedding model (unchanged)
+    # Load embedding model
     logger.info("Loading SentenceTransformer model")
     embedder = SentenceTransformer(sentence_transformer_model, device="cpu")
     logger.info("SentenceTransformer model loaded")
 
-    # Compute FAQ embeddings (unchanged)
+    # Compute FAQ embeddings
     faq_questions = [faq["question"] for faq in faqs]
     faq_embeddings = embedder.encode(faq_questions, convert_to_numpy=True).astype("float32")
     faiss.normalize_L2(faq_embeddings)
 
-    # Load Cogito model
+    # Load the 8B Cogito model
     logger.info(f"Loading {filename} model")
     model_path = hf_hub_download(
         repo_id=repo_id,
@@ -73,11 +74,12 @@ try:
         local_dir="/app/cache" if os.getenv("HF_HOME") else None,
         token=hf_token,
     )
+    # Lower n_batch for more frequent token streaming.
     generator = Llama(
         model_path=model_path,
         n_ctx=2048,
         n_threads=2,
-        n_batch=512,
+        n_batch=128,  # Adjusted for lower latency on streaming responses
         n_gpu_layers=0,
         verbose=True,
     )
@@ -97,6 +99,7 @@ def retrieve_context(query, top_k=2):
     except Exception as e:
         logger.error(f"Error in retrieve_context: {str(e)}")
         raise
+
 # Load the full CV at startup
 with open("cv_text.txt", "r", encoding="utf-8") as f:
     full_cv_text = f.read()
@@ -136,7 +139,7 @@ def stream_response(query):
         {"role": "user", "content": query}
     ]
     
-    buffer = ""
+    # Stream tokens immediately as they are generated, avoiding additional buffering.
     for chunk in generator.create_chat_completion(
         messages=messages,
         max_tokens=512,
@@ -145,22 +148,14 @@ def stream_response(query):
         top_p=0.7,
         repeat_penalty=1.2
     ):
-        text = chunk['choices'][0]['delta'].get('content', '')
-        if text:
-            buffer += text
-            if not first_token_logged and time.time() - start_time > 0:
+        token = chunk['choices'][0]['delta'].get('content', '')
+        if token:
+            if not first_token_logged:
                 logger.info(f"First token time: {time.time() - start_time:.2f}s")
                 first_token_logged = True
-                
-            # More natural chunking - yield complete sentences when possible
-            if any(buffer.endswith(char) for char in [".", "!", "?"]) or len(buffer) > 30:
-                yield f"data: {buffer}\n\n"
-                buffer = ""
-    if buffer:  # Flush remaining buffer
-        yield f"data: {buffer}\n\n"
+            yield f"data: {token}\n\n"
     yield "data: [DONE]\n\n"
 
-
 class QueryRequest(BaseModel):
     data: list
 
@@ -178,21 +173,18 @@ async def health_check():
 @app.get("/model_info")
 async def model_info():
     return {
-        "model_name": "deepcogito_cogito-v1-preview-llama-3B-GGUF",
-        "model_size": "3B",
-        "quantization": "Q4_K_M",
+        "model_name": "deepcogito_cogito-v1-preview-llama-8B-GGUF",
+        "model_size": "8B",
+        "quantization": "Q4_KM",
         "embedding_model": sentence_transformer_model,
         "faiss_index_size": len(cv_chunks),
         "faiss_index_dim": cv_embeddings.shape[1],
     }
 
-# Optimize the model loading process
-
-# Use a smaller warm-up query
+# Use a smaller warm-up query to prime the model without extensive delay.
 @app.on_event("startup")
 async def warm_up_model():
     logger.info("Warming up the model...")
-    dummy_query = "Hello"  # Shorter query
-    # Just execute once to prime the model without waiting for completion
+    dummy_query = "Hello"
     next(stream_response(dummy_query))
-    logger.info("Model warm-up initiated.")
+    logger.info("Model warm-up initiated.")