KillerKing93/Transformers-TextEngine-OpenAPI

1#!/usr/bin/env python32"""3Simple test script to verify infer_stream method works correctly4"""5 6import sys7import os8import time9 10# Add current directory to path11sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))12 13from main import Engine14 15def test_infer_stream():16    """Test the infer_stream method with a simple prompt"""17    print("=== Testing infer_stream method ===")18 19    # Initialize engine20    print("1. Initializing engine...")21    try:22        engine = Engine()23        print(f"   Engine initialized successfully with model: {engine.model_id}")24    except Exception as e:25        print(f"   Error initializing engine: {e}")26        import traceback27        traceback.print_exc()28        return False29 30    # Test messages31    messages = [32        {"role": "user", "content": "Hello, can you count to 3?"}33    ]34 35    print("\n2. Starting streaming test...")36    print(f"   Input messages: {messages}")37 38    try:39        piece_count = 040        generated_text = ""41 42        for piece in engine.infer_stream(43            messages=messages,44            max_tokens=50,45            temperature=0.746        ):47            piece_count += 148            print(f"   Piece {piece_count}: '{piece}'")49            generated_text += piece50 51        print(f"\n3. Streaming completed!")52        print(f"   Total pieces: {piece_count}")53        print(f"   Generated text: '{generated_text}'")54 55        if piece_count == 0:56            print("   ❌ ERROR: No pieces were generated!")57            return False58        else:59            print("   ✅ SUCCESS: Streaming worked correctly!")60            return True61 62    except Exception as e:63        print(f"   ❌ ERROR during streaming: {e}")64        import traceback65        traceback.print_exc()66        return False67 68def test_infer_non_stream():69    """Test the regular infer method for comparison"""70    print("\n=== Testing infer method (non-streaming) ===")71 72    try:73        engine = Engine()74        messages = [75            {"role": "user", "content": "Hello, can you count to 3?"}76        ]77 78        print("1. Starting non-streaming inference...")79        result = engine.infer(80            messages=messages,81            max_tokens=50,82            temperature=0.783        )84 85        print(f"2. Non-streaming result: '{result}'")86        print("   ✅ Non-streaming inference works!")87        return True88 89    except Exception as e:90        print(f"   ❌ ERROR during non-streaming: {e}")91        import traceback92        traceback.print_exc()93        return False94 95if __name__ == "__main__":96    print("Testing inference methods...")97    print("=" * 50)98 99    # Test non-streaming first (should work)100    non_stream_ok = test_infer_non_stream()101 102    # Test streaming (currently problematic)103    stream_ok = test_infer_stream()104 105    print("\n" + "=" * 50)106    print("SUMMARY:")107    print(f"Non-streaming: {'✅ PASS' if non_stream_ok else '❌ FAIL'}")108    print(f"Streaming: {'✅ PASS' if stream_ok else '❌ FAIL'}")109 110    if not stream_ok:111        print("\n🔍 Debugging notes:")112        print("- Check the detailed logs above for where streaming fails")113        print("- Verify the model.generate() call in the thread")114        print("- Check if TextIteratorStreamer receives any tokens")115        print("- Look for template or tokenization issues")