Spaces:
Sleeping
Sleeping
File size: 2,022 Bytes
acd8e16 |
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 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
#!/usr/bin/env python3
"""
Quick test script to verify the system works with small models.
"""
import os
import sys
from langchain_models import langchain_models_registry
from custom_evaluator import custom_evaluator
def test_smallest_model():
"""Test with the smallest available model."""
print("π Testing with smallest model (DistilGPT-2)...")
# Get the smallest model
model_config = langchain_models_registry.get_model_config("DistilGPT-2")
if not model_config:
print("β DistilGPT-2 model not found")
return False
print(f"π Model: {model_config.name}")
print(f"π Model ID: {model_config.model_id}")
try:
# Create the model
print("π₯ Creating model...")
model = langchain_models_registry.create_langchain_model(model_config)
print("β
Model created successfully")
# Test SQL generation
print("π Testing SQL generation...")
prompt_template = """
You are an expert SQL developer.
Database Schema:
{schema}
Question: {question}
Generate a SQL query:
"""
schema = "-- NYC Taxi Dataset\nCREATE TABLE trips (id INT, fare_amount FLOAT, total_amount FLOAT);"
question = "How many trips are there?"
result = langchain_models_registry.generate_sql(
model_config, prompt_template, schema, question
)
print(f"π Generated SQL: {result}")
if result and len(result) > 10:
print("β
SQL generation successful!")
return True
else:
print("β οΈ SQL generation produced short result")
return False
except Exception as e:
print(f"β Error: {e}")
return False
if __name__ == "__main__":
success = test_smallest_model()
if success:
print("\nπ System is working! Ready to run full evaluation.")
else:
print("\nβ System needs fixes.")
sys.exit(0 if success else 1)
|