Update src/app.py

added access to custom embedding and small language models

src/app.py

@@ -12,12 +12,14 @@ import doc_loader
 import modules.admin_panel as admin_panel
 
 from openai import OpenAI
-from google import genai # NEW: Google SDK
-from google.genai import types # NEW: Types for config
+from google import genai
+from google.genai import types
 from datetime import datetime
 from test_integration import run_tests
 from core.QuizEngine import QuizEngine
 from core.PineconeManager import PineconeManager
+from huggingface_hub import hf_hub_download
+from llama_cpp import Llama
 
 # --- CONFIGURATION ---
 st.set_page_config(page_title="Navy AI Toolkit", page_icon="⚓", layout="wide")
@@ -156,7 +158,52 @@ def query_model_universal(messages, max_tokens, model_choice, user_key=None):
         except Exception as e:
             return f"[OpenAI Error: {e}]", None
 
-    # --- ROUTE 3: LOCAL/OPEN SOURCE ---
+    # --- ROUTE 3: CUSTOM LOCAL GGUF (Gemma 2) ---
+    elif "Custom Gemma" in model_choice:
+        try:
+            # 1. Download Model (Cached automatically)
+            repo_id = "NavyDevilDoc/navy-custom-models" 
+            filename = "gemma-2-9b-it.Q4_K_M.gguf"
+            
+            model_path = hf_hub_download(repo_id=repo_id, filename=filename)
+
+            # 2. Initialize Llama (The Engine)
+            # n_ctx=8192 matches Gemma 2's window. n_threads=8 utilizes your CPU Upgrade.
+            llm = Llama(
+                model_path=model_path,
+                n_ctx=8192, 
+                n_threads=8, 
+                verbose=False
+            )
+
+            # 3. Format Prompt for Gemma 2 (It is picky about ChatML/Instruction format)
+            # Gemma 2 format: <start_of_turn>user\n{prompt}<end_of_turn>\n<start_of_turn>model\n
+            full_prompt = ""
+            for m in messages:
+                role = "model" if m["role"] == "assistant" else "user"
+                full_prompt += f"<start_of_turn>{role}\n{m['content']}<end_of_turn>\n"
+            full_prompt += "<start_of_turn>model\n"
+
+            # 4. Generate
+            output = llm(
+                full_prompt, 
+                max_tokens=max_tokens, 
+                stop=["<end_of_turn>"], 
+                temperature=0.3
+            )
+            
+            response_text = output['choices'][0]['text']
+            usage = {
+                "input": output['usage']['prompt_tokens'],
+                "output": output['usage']['completion_tokens']
+            }
+            return response_text, usage
+
+        except Exception as e:
+            return f"[GGUF Error: {e}]", None
+    
+    
+    # --- ROUTE 4: LOCAL/OPEN SOURCE ---
     else:
         model_map = {
             "Granite 4 (IBM)": "granite4:latest", 
@@ -264,6 +311,7 @@ with st.sidebar:
         "Standard (All-MiniLM, 384d)": "sentence-transformers/all-MiniLM-L6-v2",
         "High-Perf (MPNet, 768d)": "sentence-transformers/all-mpnet-base-v2",
         "OpenAI Small (1536d)": "text-embedding-3-small"
+        "Custom Navy (BGE, 768d)": "NavyDevilDoc/navy-custom-models/bge-finetuned" # NEW
     }
     embed_choice_label = st.selectbox("Select Embedding Model", list(embed_options.keys()))
     st.session_state.active_embed_model = embed_options[embed_choice_label]