<commit_before>from zipfile import ZipFile
from pandas import read_csv

from cassandra.cqlengine.models import Model
from cassandra.cqlengine.columns import *
from cassandra.cqlengine.connection import set_session


class Movie(Model):
    __table_name__ = 'movies'
    id = Integer(primary_key=True)
    name = Text()
    release_date = Date()
    video_release_date = Date()
    url = Text()
    tags = Set(Text)

class User(Model):
    __table_name__ = "users"
    id = Integer(primary_key=True)
    age = Integer()
    gender = Text()
    occupation = Text()
    zip = Text()


def main(context):
    context.feedback("Installing movielens")
    fp = context.download("http://files.grouplens.org/datasets/movielens/ml-100k.zip")
    zf = ZipFile(file=fp)
    tmp = zf.open("ml-100k/u.item")
    items = read_csv(tmp, sep="|", header=None, index_col=0,
                    names=[ "id", "name", "release_date", "video_release_date", "url", "unknown",
                            "Action", "Adventure", "Animation", "Children's", "Comedy", "Crime",
                            "Documentary", "Drama", "Fantasy", "Film-Noir", "Horror",  "Musical",
                            "Mystery", "Romance", "Sci-Fi", "Thriller", "War", "Western"])
    for row in items.itertuples():
        # context.feedback(row.name)
        try:
            Movie.create(id=row.Index, name=row.name.encode("utf-8"),
                         url=str(row.url))
        except Exception as e:
            print e, row
    users = zf.open("ml-100k/u.user")
    users = read_csv(users, sep="|", header=None,
                     names=["id", "age", "gender", "occupation", "zip"], index_col=0)

    for user in users.itertuples():
        try:
            User.create(id=user.Index, age=user.age, gender=user.gender,
                        occupation=user.occupation, zip=user.zip)
        except Exception as e:
            print user.id, e

    data = zf.open("ml-100k/u.data")
    # ratings = read_csv(data, )


if __name__ == "__main__":
    from cdm import install_local
    install_local("movielens-small", main)<commit_msg>post install improvements<commit_after>from zipfile import ZipFile
from pandas import read_csv

from cassandra.cqlengine.models import Model
from cassandra.cqlengine.columns import *
from cassandra.cqlengine.connection import set_session


class Movie(Model):
    __table_name__ = 'movies'
    id = Integer(primary_key=True)
    name = Text()
    release_date = Date()
    video_release_date = Date()
    url = Text()
    tags = Set(Text)

class User(Model):
    __table_name__ = "users"
    id = Integer(primary_key=True)
    age = Integer()
    gender = Text()
    occupation = Text()
    zip = Text()


def main(context):
    context.feedback("Installing movielens")
    fp = context.download("http://files.grouplens.org/datasets/movielens/ml-100k.zip")
    zf = ZipFile(file=fp)
    tmp = zf.open("ml-100k/u.item")
    items = read_csv(tmp, sep="|", header=None, index_col=0,
                    names=[ "id", "name", "release_date", "video_release_date", "url", "unknown",
                            "Action", "Adventure", "Animation", "Children's", "Comedy", "Crime",
                            "Documentary", "Drama", "Fantasy", "Film-Noir", "Horror",  "Musical",
                            "Mystery", "Romance", "Sci-Fi", "Thriller", "War", "Western"])
    for row in items.itertuples():
        # context.feedback(row.name)
        try:
            Movie.create(id=row.Index, name=row.name.encode("utf-8"),
                         url=str(row.url))
        except Exception as e:
            print e, row
    users = zf.open("ml-100k/u.user")
    users = read_csv(users, sep="|", header=None,
                     names=["id", "age", "gender", "occupation", "zip"], index_col=0)

    for user in users.itertuples():
        try:
            User.create(id=user.Index, age=user.age, gender=user.gender,
                        occupation=user.occupation, zip=user.zip)
        except Exception as e:
            print user.id, e

    # user id | item id | rating | timestamp
    data = zf.open("ml-100k/u.data")
    prepared = context.session.prepare("INSERT INTO ratings_by_movie (movie_id, user_id, rating, timestamp) VALUES (?, ?, ?, ?)")
    ratings = read_csv(data, sep="|", header=None,
                       names=["user_id", "movie_id", "rating", "timestamp"])

    for r in ratings.itertuples():
        context.session.execute(prepared, (r.movie_id, r.user_id, r.rating, r.timestamp))


if __name__ == "__main__":
    from cdm import install_local
    install_local("movielens-small", main)
{'fraction_non_alphanumeric': 0.14530099610220876, 'fraction_numerical': 0.0064963187527068, 'mean_word_length': 2.5833979829325058, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '456668', 'n_tokens_mistral': 1535, 'n_tokens_neox': 1473, 'n_words': 325}