|
|
@ -0,0 +1,22 @@ |
|
|
|
from typing import Any, Dict, Iterator, Type |
|
|
|
|
|
|
|
import pandas as pd |
|
|
|
from django.db.models.query import QuerySet |
|
|
|
|
|
|
|
from .labels import Labels |
|
|
|
from examples.models import Example |
|
|
|
|
|
|
|
|
|
|
|
class Dataset: |
|
|
|
def __init__(self, examples: QuerySet[Example], user, label_collection_class: Type[Labels], confirmed_only=False): |
|
|
|
if confirmed_only: |
|
|
|
examples = examples.filter(states__confirmed_by=user) |
|
|
|
self.examples = examples |
|
|
|
self.labels = label_collection_class(examples, user) |
|
|
|
|
|
|
|
def __iter__(self) -> Iterator[Dict[str, Any]]: |
|
|
|
for example in self.examples: |
|
|
|
yield {"id": example.id, "data": example.text, **example.meta, **self.labels.find_by(example.id)} |
|
|
|
|
|
|
|
def to_pandas(self) -> pd.DataFrame: |
|
|
|
return pd.DataFrame(self) |