Replace Gemini with in-Space open-source LLM (Qwen3.5-2B via llama.cpp)

app.py

@@ -68,6 +68,7 @@ LOCAL_LLM_FILE = os.environ.get("LOCAL_LLM_FILE", "Qwen3.5-2B-Q4_K_M.gguf")
 LOCAL_LLM_CTX = int(os.environ.get("LOCAL_LLM_CTX", "4096"))
 LOCAL_LLM_THREADS = int(os.environ.get("LOCAL_LLM_THREADS", "2"))
 LOCAL_LLM_MAX_TOKENS = int(os.environ.get("LOCAL_LLM_MAX_TOKENS", "320"))
+LOCAL_LLM_TEMPERATURE = float(os.environ.get("LOCAL_LLM_TEMPERATURE", "0.3"))
 LOCAL_LLM_MAX_INPUT_CHARS = int(os.environ.get("LOCAL_LLM_MAX_INPUT_CHARS", "2400"))
 LOCAL_LLM_TIME_BUDGET_S = float(os.environ.get("LOCAL_LLM_TIME_BUDGET_S", "22"))
 LOCAL_LLM_QUEUE_TIMEOUT_S = float(os.environ.get("LOCAL_LLM_QUEUE_TIMEOUT_S", "5"))
@@ -306,8 +307,8 @@ class MLModelLoader:
                     loaded += 1
                     logger.info(f"✅ Loaded model from Hub: {name}")
                     continue
-                except Exception:
-                    pass
+                except Exception as e:
+                    logger.warning(f"Could not load {name} from the Hub: {type(e).__name__}: {str(e)[:200]}")
 
                 # 2. Try all local search directories
                 for sdir in search_dirs:
@@ -593,30 +594,33 @@ class LocalLLM:
                     f"<|im_start|>assistant\n{self._prefill}"
                 )
                 stream = self.llm.create_completion(
-                    prompt, max_tokens=max_tokens, temperature=0.7, top_p=0.8, top_k=20,
+                    prompt, max_tokens=max_tokens, temperature=LOCAL_LLM_TEMPERATURE, top_p=0.8, top_k=20,
                     presence_penalty=1.5, stop=["<|im_end|>", "<|endoftext|>"], stream=True,
                 )
                 read = lambda c: c["choices"][0].get("text") or ""
             else:
                 stream = self.llm.create_chat_completion(
                     messages=[{"role": "system", "content": system}, {"role": "user", "content": user}],
-                    max_tokens=max_tokens, temperature=0.7, top_p=0.8, stream=True,
+                    max_tokens=max_tokens, temperature=LOCAL_LLM_TEMPERATURE, top_p=0.8, stream=True,
                 )
                 read = lambda c: c["choices"][0].get("delta", {}).get("content") or ""
             pieces: List[str] = []
             truncated = False
             for chunk in stream:
                 pieces.append(read(chunk))
+                if chunk["choices"][0].get("finish_reason") == "length":
+                    truncated = True
                 if time.monotonic() - started > LOCAL_LLM_TIME_BUDGET_S:
                     truncated = True
                     break
             import re
             text = re.sub(r"<think>.*?</think>", "", "".join(pieces), flags=re.S).strip()
             if truncated:
-                # Out of time: end on the last complete sentence or line
-                cut = max(text.rfind(". "), text.rfind(".\n"), text.rfind("\n"))
-                if cut > len(text) * 0.5:
-                    text = text[:cut + 1].strip()
+                # Out of time or tokens: end on the last complete sentence. A dot
+                # after a digit is a list marker ("2."), not a sentence end.
+                ends = [m.end() for m in re.finditer(r"(?<!\d)[.!?][*_`\")\]]*(?=\s|$)", text)]
+                if ends and ends[-1] > len(text) * 0.4:
+                    text = text[:ends[-1]].strip()
             if not text:
                 return {**base, "response": None, "error": "local LLM returned no text"}
             return {
@@ -811,6 +815,16 @@ class TransformerModelLoader:
             logger.warning(f"Hosted LLM error: {type(e).__name__}: {str(e)[:160]} — using the local model")
         return None
 
+    # The in-Space model is small and slow: ask for a short answer it can finish
+    # inside the time budget, and tell it not to fill gaps with invented detail.
+    LOCAL_SYSTEM_PROMPT = (
+        "You are CyberForge AI, a cybersecurity analyst. Answer in at most 100 words, "
+        "as short bullet points. When assessing a threat, start with "
+        "'Risk Level: Critical, High, Medium or Low'. State only what you are sure of: "
+        "never invent CVE numbers, MITRE ATT&CK IDs, commands or reputation data, and "
+        "say so plainly when you do not know."
+    )
+
     def security_chat(self, query: str, max_tokens: int = 512, allow_local: bool = True) -> Dict:
         """Cybersecurity Q&A.
         1. Hosted model through HF Inference Providers, while the account has credit.
@@ -822,7 +836,7 @@ class TransformerModelLoader:
                 return hosted
         if not allow_local:
             return {"model": "security-llm", "response": None, "error": "hosted LLM unavailable"}
-        result = local_llm.chat(self.SYSTEM_PROMPT, query, max_tokens=max_tokens)
+        result = local_llm.chat(self.LOCAL_SYSTEM_PROMPT, query, max_tokens=max_tokens)
         if result.get("response"):
             self.last_engine = "local"
         return result

requirements.txt

@@ -2,7 +2,9 @@
 # Combined Gradio UI + FastAPI REST API
 
 # Core ML
-scikit-learn>=1.5.0
+# Pinned: the .pkl files in Che237/cyberforge-models were saved with 1.8.0 and the
+# gradient-boosting ones do not load on 1.9+. Retrain and re-upload before raising it.
+scikit-learn==1.8.0
 pandas>=2.2.0
 numpy>=2.0.0
 joblib>=1.4.0