Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,182 Bytes
14ffd10 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
#!/usr/bin/env python3
"""
Test script to verify thinking extraction works correctly
"""
import re
from typing import Tuple
def extract_thinking_and_response(text: str) -> Tuple[str, str]:
"""Extract thinking process and clean response from AI output"""
thinking = ""
response = text
# Extract thinking content
thinking_match = re.search(r'<think>(.*?)</think>', text, re.DOTALL)
if thinking_match:
thinking = thinking_match.group(1).strip()
response = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)
# Clean up the response
response = re.sub(r'^(Assistant:|AI:|Response:|Answer:)\s*', '', response.strip())
response = re.sub(r'\[INST\].*?\[\/INST\]', '', response)
response = re.sub(r'<\|.*?\|>', '', response)
return thinking.strip(), response.strip()
# Test cases
test_cases = [
# Basic thinking extraction
"<think>I need to analyze this query and generate search terms.</think>Here are the search queries:\n1. AI developments 2024\n2. artificial intelligence news",
# No thinking
"Here are the search queries:\n1. AI developments 2024\n2. artificial intelligence news",
# Complex thinking
"<think>The user is asking about quantum computing breakthroughs. I should focus on recent developments, key research areas, and practical applications. Let me think about the best search terms...</think>Based on your question, here are optimized search queries for quantum computing breakthroughs.",
# Thinking with newlines
"<think>\nThis is a complex question about climate change.\nI need to consider multiple aspects:\n1. Ocean currents\n2. Temperature changes\n3. Recent research\n</think>Climate change affects ocean currents in several ways..."
]
print("Testing thinking extraction...")
for i, test in enumerate(test_cases, 1):
thinking, response = extract_thinking_and_response(test)
print(f"\nTest {i}:")
print(f"Original: {test[:100]}...")
print(f"Thinking: {thinking[:100]}..." if thinking else "Thinking: None")
print(f"Response: {response[:100]}...")
print("-" * 50)
print("✅ Thinking extraction test completed!") |