File size: 2,224 Bytes
e57b93d b0d65d0 1cd5388 d90b23a 176b07f b0d65d0 9499c2d 176b07f 9499c2d 176b07f b0d65d0 176b07f 35358b3 176b07f e57b93d 176b07f 35358b3 176b07f ffd7b2d 35358b3 176b07f 1cd5388 b0d65d0 | 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 71 72 73 | from lida.components import Manager
from llmx import llm, TextGenerationConfig
import os
lida = Manager(text_gen=llm("openai"))
cars_data_url = "https://raw.githubusercontent.com/uwdata/draco/master/data/cars.csv"
def test_summarizer():
textgen_config = TextGenerationConfig(
n=1, temperature=0, use_cache=False, max_tokens=None)
summary_no_enrich = lida.summarize(
cars_data_url,
textgen_config=textgen_config,
summary_method="default")
summary_enrich = lida.summarize(cars_data_url,
textgen_config=textgen_config, summary_method="llm")
assert summary_no_enrich != summary_enrich
assert "dataset_description" in summary_enrich and len(
summary_enrich["dataset_description"]) > 0
def test_goals():
textgen_config = TextGenerationConfig(
n=1, temperature=0.1, use_cache=False, max_tokens=None)
summary = lida.summarize(
cars_data_url,
textgen_config=textgen_config, summary_method="default")
goals = lida.goals(summary, n=2, textgen_config=textgen_config)
assert len(goals) == 2
assert len(goals[0].question) > 0
def test_vizgen():
textgen_config = TextGenerationConfig(
n=1,
temperature=0.1,
use_cache=True,
max_tokens=None)
summary = lida.summarize(
cars_data_url,
textgen_config=textgen_config, summary_method="default")
goals = lida.goals(summary, n=2, textgen_config=textgen_config)
charts = lida.visualize(
summary=summary,
goal=goals[0],
textgen_config=textgen_config,
library="seaborn")
assert len(charts) > 0
first_chart = charts[0]
# Ensure the first chart has a status of True
assert first_chart.status is True
# Ensure no errors in the first chart
assert first_chart.error is None
# Ensure the raster image of the first chart exists
assert len(first_chart.raster) > 0
# Test saving the raster image of the first chart
temp_file_path = "temp_image.png"
first_chart.savefig(temp_file_path)
# Ensure the image is saved correctly
assert os.path.exists(temp_file_path)
# Clean up
os.remove(temp_file_path)
|