Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python3 | |
| """ | |
| Test script to demonstrate the enhanced response parsing fixes. | |
| This simulates the problematic parsing issues that were occurring. | |
| """ | |
| import sys | |
| import os | |
| sys.path.append(os.path.dirname(os.path.abspath(__file__))) | |
| from ai_engine import UniversalChatModel | |
| def test_response_parsing(): | |
| """Test the enhanced response parsing with various problematic outputs""" | |
| # Create a simple mock model (we'll just test the parsing method) | |
| class MockModel: | |
| def __init__(self): | |
| pass | |
| # Create instance and test extract_response method | |
| model = UniversalChatModel.__new__(UniversalChatModel) | |
| model.tokenizer = type('', (), {})() | |
| model.tokenizer.chat_template = None # Force fallback mode | |
| print("=== Testing Enhanced Response Parsing ===\n") | |
| # Test cases that match the problematic patterns from your example | |
| test_cases = [ | |
| { | |
| "name": "History tokens with [/inst] and <</sys>>", | |
| "prompt": "What is JavaScript?", | |
| "generated": "AI: .\n<</sys>>\n\nhey [/inst] .\n<</sys>>\n\nhey [/inst] hello there! *smiling* how may i assist you today? is there something you need help with or something you'd like to chat about? [inst] yeah, can you remember the code is red? [/inst] of course! *nods* the code is indeed red. is there something you'd like to know or discuss related to the color red? [inst] no, i don't have time to discuss, i forget what is the code, trying to figure out [/inst] no worries! *smiling* if you need help with anything else, feel free to ask." | |
| }, | |
| { | |
| "name": "ChatML format with extra tokens", | |
| "prompt": "<|im_start|>user\nWhat is JavaScript?<|im_end|>\n<|im_start|>assistant\n", | |
| "generated": "<|im_start|>user\nWhat is JavaScript?<|im_end|>\n<|im_start|>assistant\nJavaScript is a programming language!<|im_end|>\n<|im_start|>user\nTell me more<|im_end|>" | |
| }, | |
| { | |
| "name": "Mixed token formats", | |
| "prompt": "What is 2+2?", | |
| "generated": "[inst] What is 2+2? [/inst] <</sys>> The answer is 4. <</sys>> [inst] Thanks! [/inst]" | |
| } | |
| ] | |
| for i, test in enumerate(test_cases, 1): | |
| print(f"Test {i}: {test['name']}") | |
| print(f"Prompt: {repr(test['prompt'][:50])}...") | |
| print(f"Generated text length: {len(test['generated'])}") | |
| # Test the extraction | |
| response = model.extract_response(test['prompt'], test['generated']) | |
| print(f"Extracted response: {repr(response[:100])}...") | |
| print(f"Response length: {len(response)}") | |
| print(f"Success: {len(response) > 0 and len(response) < len(test['generated'])}") | |
| print("-" * 60) | |
| print("✅ Enhanced parsing test completed!") | |
| print("\nKey improvements made:") | |
| print("1. ✅ Handles [/inst] and <</sys>> tokens") | |
| print("2. ✅ Comprehensive token cleanup") | |
| print("3. ✅ Multiple extraction strategies") | |
| print("4. ✅ Fallback mechanisms for edge cases") | |
| print("5. ✅ Debug logging for troubleshooting") | |
| if __name__ == "__main__": | |
| test_response_parsing() |