Update app.py

app.py

@@ -5,6 +5,7 @@ from tokenizers import Tokenizer
 from huggingface_hub import hf_hub_download
 from ModelArchitecture import Transformer, ModelConfig, generate
 from safetensors.torch import load_file
+import re
 
 # -----------------------------
 # Load model and tokenizer
@@ -12,7 +13,7 @@ from safetensors.torch import load_file
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 REPO_ID = "VirtualInsight/Lumen-Instruct"
 
-# Download model assets from Hugging Face Hub
+# Download model assets
 model_path = hf_hub_download(repo_id=REPO_ID, filename="model.safetensors")
 tokenizer_path = hf_hub_download(repo_id=REPO_ID, filename="tokenizer.json")
 config_path = hf_hub_download(repo_id=REPO_ID, filename="config.json")
@@ -27,7 +28,7 @@ model.load_state_dict(load_file(model_path, device=str(device)), strict=False)
 model.eval()
 
 # -----------------------------
-# Special Tokens for Chat Format
+# Special Tokens
 # -----------------------------
 EOS_TOKEN = "<|im_end|>"
 EOS_TOKEN_ID = tokenizer.encode(EOS_TOKEN).ids[0]
@@ -41,13 +42,13 @@ def generate_response(prompt, max_tokens=200, temperature=0.7, top_p=0.9):
     """
     Generates a clean assistant-only response from the Lumen Instruct model.
     """
-    # Format input as a conversation prompt
+    # Chat-style input
     formatted_prompt = f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
 
     # Tokenize input
     input_ids = torch.tensor([tokenizer.encode(formatted_prompt).ids], dtype=torch.long, device=device)
 
-    # Generate output
+    # Generate
     output = generate(
         model,
         input_ids,
@@ -59,18 +60,27 @@ def generate_response(prompt, max_tokens=200, temperature=0.7, top_p=0.9):
         eos_token_id=EOS_TOKEN_ID,
     )
 
-    # Decode full text
+    # Decode text
     full_text = tokenizer.decode(output[0].tolist())
 
-    # 🧹 Clean extraction of assistant’s reply only
+    # -----------------------------
+    # Clean assistant-only response
+    # -----------------------------
+    # 1. Get part after last assistant marker
     if "<|im_start|>assistant" in full_text:
         response = full_text.split("<|im_start|>assistant")[-1]
-        response = response.split("<|im_end|>")[0] if "<|im_end|>" in response else response
     else:
         response = full_text
 
-    # Remove potential leftover role tokens and clean spaces
-    response = response.replace("assistant", "").replace("user", "").strip()
+    # 2. Cut off at end marker if exists
+    response = response.split("<|im_end|>")[0]
+
+    # 3. Remove any lingering user/assistant labels or context lines
+    response = re.sub(r"(?i)\buser\b.*", "", response)
+    response = re.sub(r"(?i)\bassistant\b.*", "", response)
+
+    # 4. Clean newlines and whitespace
+    response = response.strip()
 
     return response
 
@@ -91,7 +101,7 @@ demo = gr.Interface(
 )
 
 # -----------------------------
-# Launch Interface
+# Launch
 # -----------------------------
 if __name__ == "__main__":
     demo.launch(share=True)