Update app.py

app.py

@@ -7,6 +7,7 @@ import shutil
 import torch
 import psutil
 import time
+from orbit import orbit
 
 # Define path for HF cache to clean
 HF_CACHE_DIR = os.path.expanduser("~/.cache/huggingface/hub")
@@ -25,43 +26,44 @@ MODELS = [
     'MihaiPopa-1/CinnabarLM-1.5M-Base', 'Harley-ml/Dillionv2-1.3M', 'Eclipse-Senpai/KeyLM-75M',
     'SupraLabs/Supra-Mini-v6-1M', 'AxiomicLabs/GPT-S-1.4M', 'GODELEV/Archaea-74M',
     'Sandroeth/cali-0.1B', 'veyra-ai/veyra3-5m-base', 'veyra-ai/veyra-30m-base-5b-tokens',
-    'ThingAI/Quark-50m', 'ThingAI/Quark-135m', 'HuggingFaceTB/SmolLM2-135M-Instruct', 
+    'ThingAI/Quark-50m', 'ThingAI/Quark-135m', 'HuggingFaceTB/SmolLM2-135M-Instruct',
     'Aravindan/awesome-gpt-2-coder', 'Qwen/Qwen2.5-Coder-0.5B', 'SupraLabs/Supra-50M-Reasoning'
 ]
 
 ACTIVE_SESSIONS = {}
 SESSION_TIMEOUT = 60
 
+
 def live_count(request: gr.Request):
     current_time = time.time()
     if request:
         ACTIVE_SESSIONS[request.session_hash] = current_time
-        
-    # Prune
+
     expired = [s for s, t in ACTIVE_SESSIONS.items() if current_time - t > SESSION_TIMEOUT]
     for s in expired:
         ACTIVE_SESSIONS.pop(s, None)
-        
+
     return len(ACTIVE_SESSIONS)
 
-# Global class to safely manage the loaded model and tokenizer in memory
+
 class ModelManager:
     def __init__(self):
         self.model = None
         self.tokenizer = None
         self.model_id = None
-        self.stop_generation = False # Added flag to instantly kill generation
+        self.stop_generation = False
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
 
+
 model_manager = ModelManager()
 
-# Custom stopping criteria to halt the generation thread when loading a new model
+
 class StopOnFlag(StoppingCriteria):
     def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
         return model_manager.stop_generation
 
+
 def get_system_stats(request: gr.Request = None):
-    """Returns a dictionary of current system metrics with formatted strings."""
     mem = psutil.virtual_memory()
     disk = psutil.disk_usage('/')
     return (
@@ -71,50 +73,62 @@ def get_system_stats(request: gr.Request = None):
         f"Active\t: \t{len(ACTIVE_SESSIONS) if request is None else live_count(request)} session(s)"
     )
 
+
 def load_new_model(model_id):
-    """Loads the model and tokenizer dynamically into the global manager."""
-    # Stop any ongoing generation immediately
     model_manager.stop_generation = True
-    
-    # Clear old model from memory
+
     model_manager.model = None
     model_manager.tokenizer = None
     model_manager.model_id = None
+
     yield f"Loading {model_id}..."
+
     gc.collect()
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
-    
+
     try:
-        # Load explicitly for streaming purposes instead of pipeline
         tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
         model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True).to(model_manager.device)
-        
+
         model_manager.tokenizer = tokenizer
         model_manager.model = model
         model_manager.model_id = model_id
-        
+
         yield f"Successfully loaded {model_id} on {model_manager.device.upper()}"
+
     except Exception as e:
         yield f"Error loading model: {str(e)}"
 
+
 def run_inference(user_prompt, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample):
-    """Generates text via streaming generator."""
     if model_manager.model is None or model_manager.tokenizer is None:
         yield "Please load a model first.", "Model not loaded"
         return
-        
-    # Reset the stop flag for the new generation run
+
     model_manager.stop_generation = False
-    
+
     tokenizer = model_manager.tokenizer
     model = model_manager.model
     model_id = model_manager.model_id
-    
+
+    # ORBIT GOVERNOR INTERCEPT
+    original_prompt = user_prompt
+    decision = orbit.as_dict(user_prompt)
+
+    if not decision["allow"]:
+        yield decision["message"], "Orbit blocked prompt"
+        return
+
+    user_prompt = decision["prompt"]
+
     is_supra_reasoning = "Supra-50M-Reasoning" in model_id if model_id else False
-    
+
     if is_supra_reasoning:
-        SYSTEM_PROMPT = "Your role as an assistant involves thoroughly exploring questions through a systematic long thinking process before providing the final precise and accurate solutions."
+        SYSTEM_PROMPT = (
+            "Your role as an assistant involves thoroughly exploring questions through "
+            "a systematic long thinking process before providing the final precise and accurate solutions."
+        )
         prompt_to_encode = (
             f"[SYSTEM]: {SYSTEM_PROMPT}\n\n"
             f"[USER]: {user_prompt}\n\n"
@@ -125,17 +139,18 @@ def run_inference(user_prompt, max_tokens, temperature, top_k, top_p, rep_penalt
         prompt_to_encode = user_prompt
         skip_special = True
 
-    # Tokenize input
     inputs = tokenizer([prompt_to_encode], return_tensors="pt").to(model_manager.device)
-    
-    # Set up the streamer
-    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=skip_special)
-    
-    # Adjust variables based on the do_sample logic
+
+    streamer = TextIteratorStreamer(
+        tokenizer,
+        timeout=60.0,
+        skip_prompt=True,
+        skip_special_tokens=skip_special
+    )
+
     if not do_sample:
-        temperature = 1.0 # Temperature is ignored if do_sample=False, but setting it > 0 avoids config errors
+        temperature = 1.0
 
-    # Generation arguments
     generate_kwargs = dict(
         **inputs,
         streamer=streamer,
@@ -146,42 +161,41 @@ def run_inference(user_prompt, max_tokens, temperature, top_k, top_p, rep_penalt
         repetition_penalty=float(rep_penalty),
         no_repeat_ngram_size=int(ngram_size),
         do_sample=do_sample,
-        pad_token_id=tokenizer.eos_token_id, # Prevents padding warnings
-        stopping_criteria=StoppingCriteriaList([StopOnFlag()]) # Attach the stopping criteria
+        pad_token_id=tokenizer.eos_token_id,
+        stopping_criteria=StoppingCriteriaList([StopOnFlag()])
     )
 
     start_time = time.time()
-    # Start generation in a separate background thread
     thread = Thread(target=model.generate, kwargs=generate_kwargs)
     thread.start()
-    
+
     if is_supra_reasoning:
-        # Use plain text formatting rather than markdown symbols inside gr.Textbox
-        base_display = f"Prompt: {user_prompt}\n\n----------------------------------------\n\n"
+        base_display = f"Prompt: {original_prompt}\n\n----------------------------------------\n\n"
         generated_text = ""
     else:
         base_display = ""
-        generated_text = user_prompt
+        generated_text = ""
 
-    # Yield output iteratively for the streaming effect
     token_count = 0
+
     for new_text in streamer:
-        # Immediately break out of the UI update loop if a new model is loaded
         if model_manager.stop_generation:
             break
-            
+
         generated_text += new_text
         token_count += 1
+
         duration = time.time() - start_time
         tps = token_count / duration if duration > 0 else 0
-        
+
         display_text = generated_text
-        
+
         if is_supra_reasoning:
             display_text = display_text.replace("<s>", "").replace("</s>", "")
+
             if not display_text.startswith("🧠 Thinking Process:"):
                 display_text = "🧠 Thinking Process:\n" + display_text
-                
+
             display_text = display_text.replace("<|begin_of_thought|>", "🧠 Thinking Process:\n")
             display_text = display_text.replace("<|end_of_thought|>", "\n\n")
             display_text = display_text.replace("<|begin_of_solution|>", "✅ Final Answer:\n\n")
@@ -189,6 +203,7 @@ def run_inference(user_prompt, max_tokens, temperature, top_k, top_p, rep_penalt
 
         yield base_display + display_text, f"Speed: {tps:.2f} tokens/sec"
 
+
 def clean_cache():
     if os.path.exists(HF_CACHE_DIR):
         shutil.rmtree(HF_CACHE_DIR)
@@ -196,73 +211,126 @@ def clean_cache():
         return "Cache cleaned successfully!"
     return "Cache directory not found."
 
-# Gradio Interface
-with gr.Blocks(title="Small MF Model Tester", theme=gr.themes.Soft()) as app:
-    
-    gr.Markdown("# 🚀 Small Model Evaluation Hub with Streaming")
+
+with gr.Blocks(title="Small MF Model Tester + Orbit Governor", theme=gr.themes.Soft()) as app:
+
+    gr.Markdown("# 🚀 Small Model Evaluation Hub with Orbit Governor")
 
     with gr.Row():
-        # Left column: Settings & Monitoring
         with gr.Column(scale=1):
-            
+
             with gr.Accordion("System Monitoring", open=True):
                 stats_output = gr.Textbox(label="Live System Stats", show_label=False)
                 gr.Timer(2).tick(get_system_stats, None, stats_output)
 
             with gr.Group():
-                gr.Markdown("### Select or Paste custom model id here")
+                gr.Markdown("### Select or paste custom model ID here")
                 with gr.Row():
-                    model_id_input = gr.Dropdown(choices=MODELS, label="Model", allow_custom_value=True, show_label=False, scale=3)
+                    model_id_input = gr.Dropdown(
+                        choices=MODELS,
+                        label="Model",
+                        allow_custom_value=True,
+                        show_label=False,
+                        scale=3
+                    )
                     load_btn = gr.Button("Load", variant="secondary", scale=1)
-                
+
                 clean_btn = gr.Button("Clean HF Cache", variant="stop", size="sm")
 
             with gr.Accordion("Generation Configuration", open=False):
-                do_sample_input = gr.Checkbox(label="Enable Sampling (do_sample)", value=True, info="Uncheck for greedy decoding")
-                max_tokens_input = gr.Slider(minimum=10, maximum=2048, value=256, step=1, label="Max Output Tokens")
-                temperature_input = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature", info="Higher = more creative")
-                
-                top_k_input = gr.Slider(minimum=0, maximum=100, value=50, step=1, label="Top-K", info="0 = disabled")
-                top_p_input = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top-P (Nucleus)", info="1.0 = disabled")
-                rep_penalty_input = gr.Slider(minimum=1.0, maximum=2.0, value=1.1, step=0.05, label="Repetition Penalty", info="1.0 = disabled")
-                ngram_size_input = gr.Slider(minimum=0, maximum=10, value=0, step=1, label="No Repeat N-Gram Size", info="0 = disabled")
-
-        # Right column: Interaction
+                do_sample_input = gr.Checkbox(
+                    label="Enable Sampling (do_sample)",
+                    value=True,
+                    info="Uncheck for greedy decoding"
+                )
+                max_tokens_input = gr.Slider(
+                    minimum=10,
+                    maximum=2048,
+                    value=256,
+                    step=1,
+                    label="Max Output Tokens"
+                )
+                temperature_input = gr.Slider(
+                    minimum=0.1,
+                    maximum=2.0,
+                    value=0.7,
+                    step=0.1,
+                    label="Temperature",
+                    info="Higher = more creative"
+                )
+                top_k_input = gr.Slider(
+                    minimum=0,
+                    maximum=100,
+                    value=50,
+                    step=1,
+                    label="Top-K",
+                    info="0 = disabled"
+                )
+                top_p_input = gr.Slider(
+                    minimum=0.1,
+                    maximum=1.0,
+                    value=0.9,
+                    step=0.05,
+                    label="Top-P (Nucleus)",
+                    info="1.0 = disabled"
+                )
+                rep_penalty_input = gr.Slider(
+                    minimum=1.0,
+                    maximum=2.0,
+                    value=1.1,
+                    step=0.05,
+                    label="Repetition Penalty",
+                    info="1.0 = disabled"
+                )
+                ngram_size_input = gr.Slider(
+                    minimum=0,
+                    maximum=10,
+                    value=0,
+                    step=1,
+                    label="No Repeat N-Gram Size",
+                    info="0 = disabled"
+                )
+
         with gr.Column(scale=2):
             user_prompt = gr.Textbox(
-                label="Prompt", 
-                value="Once upon a time in a digital kingdom,", 
-                placeholder="Enter your prompt here...", 
+                label="Prompt",
+                value="Once upon a time in a digital kingdom,",
+                placeholder="Enter your prompt here...",
                 lines=5
             )
+
             run_btn = gr.Button("Generate text", variant="primary", size="lg")
             status_output = gr.Markdown("Status: *Waiting to load model...*")
-            output_text = gr.Textbox(label="Result", lines=15, buttons=["copy"], autoscroll=True)
-    
-    # Events
+            output_text = gr.Textbox(
+                label="Result",
+                lines=15,
+                buttons=["copy"],
+                autoscroll=True
+            )
+
     load_btn.click(
-        fn=load_new_model, 
-        inputs=[model_id_input], 
+        fn=load_new_model,
+        inputs=[model_id_input],
         outputs=[status_output]
     )
-    
-    # We use `.click` targeting a generator function, which Gradio naturally treats as a streaming output
+
     run_btn.click(
         fn=run_inference,
         inputs=[
-            user_prompt, 
-            max_tokens_input, 
-            temperature_input, 
-            top_k_input, 
-            top_p_input, 
-            rep_penalty_input, 
+            user_prompt,
+            max_tokens_input,
+            temperature_input,
+            top_k_input,
+            top_p_input,
+            rep_penalty_input,
             ngram_size_input,
             do_sample_input
         ],
         outputs=[output_text, status_output]
     )
-    
+
     clean_btn.click(fn=clean_cache, outputs=[status_output])
 
+
 if __name__ == "__main__":
     app.launch()