#!/bin/python3
from plan2scene.common.house_parser import parse_houses, load_house_texture_embeddings, save_house_crops, save_house_texture_embeddings, \
    load_houses_with_embeddings
from plan2scene.common.image_description import ImageSource
from plan2scene.common.residence import House, Room
from plan2scene.config_manager import ConfigManager
import os.path as osp
import os
import logging
import torch

from plan2scene.crop_select.util import fill_textures
from plan2scene.texture_gen.predictor import TextureGenPredictor
from plan2scene.texture_gen.utils.io import load_conf_eval

def compute_rs_mean_embeddings(conf: ConfigManager, train_houses):
    surface_type_rt_embs = {surface: {} for surface in conf.surfaces}  # Mapping surface_type -> room_type -> list of embeddings

    for i, (house_key, house) in enumerate(train_houses.items()):
        assert isinstance(house, House)
        logging.info("[%d/%d] Computing RS Embedding %s" % (i, len(train_houses), house_key))

        for room_index, room in house.rooms.items():
            assert isinstance(room, Room)
            # Calculate the mean embs
            for room_type in room.types:
                for surface in conf.surfaces:
                    if len(room.surface_embeddings[surface]) > 0:
                        if room_type not in surface_type_rt_embs[surface]:
                            surface_type_rt_embs[surface][room_type] = []
                        surface_type_rt_embs[surface][room_type].extend([a for a in room.surface_embeddings[surface].values()])

    surface_rt_emb_map = {surface: {} for surface in conf.surfaces}  # Mapping surface_type -> room_type -> mean embedding
    for surface in conf.surfaces:
        for room_type in surface_type_rt_embs[surface]:
            surface_rt_emb_map[surface][room_type] = torch.mean(torch.cat(surface_type_rt_embs[surface][room_type]), dim=0).unsqueeze(0)

    return surface_rt_emb_map

def process_house(conf: ConfigManager, house: House, predictor: TextureGenPredictor, rs_embeddings: dict):
    for room_index, room in house.rooms.items():
        assert isinstance(room, Room)
        for surface in conf.surfaces:
            if "prop" in room.surface_embeddings[surface]:
                continue  # Already has a prediction. Surface observed. Therefore, skip.

            # Collect candidates from different room type labels assigned to room
            surface_emb_candidates = []
            for room_type in room.types:
                if room_type in rs_embeddings[surface]:
                    surface_emb_candidates.append(rs_embeddings[surface][room_type])

            if len(surface_emb_candidates) > 0:
                surface_emb = torch.cat(surface_emb_candidates, dim=0).mean(dim=0).unsqueeze(0)
                room.surface_embeddings[surface] = {"prop": surface_emb}

    fill_textures(conf, {house.house_key: house}, predictor=predictor, log=False, image_source=ImageSource.RS_MEAN_EMB, skip_existing_textures=False)

if __name__ == "__main__":
    """
    Predict textures for unobserved surfaces using the NaiveSynth approach.
    """

    import argparse

    conf = ConfigManager()
    parser = argparse.ArgumentParser(description="Predict textures for unobserved surfaces using the NaiveSynth approach.")
    conf.add_args(parser)
    parser.add_argument("output_path",
                        help="Output path to save the predicted embeddings and crops."
                             "Usually './data/processed/baselines/naivesynth/all_surfaces/[split]/drop_[drop_fraction]'.")
    parser.add_argument("observed_surfaces_predictions_path", help="Path to saved texture_gen embeddings and crops for observed surfaces. "
                                                                   "Usually './data/processed/baselines/naivesynth/observed/[split]/drop_[drop_fraction]'.")
    parser.add_argument("split", help="val/test")
    parser.add_argument("train_embeddings_path",
                        help="Path to saved texture_gen embeddings for the train set. "
                             "Usually './data/processed/baselines/naivesynth/texture_gen/train/drop_0.0/surface_texture_embeddings'")
    args = parser.parse_args()
    conf.process_args(args, output_is_dir=True)

    output_path = args.output_path
    split = args.split
    observed_surfaces_predictions_path = args.observed_surfaces_predictions_path
    train_embeddings_path = args.train_embeddings_path

    if not osp.exists(osp.join(output_path, "texture_crops")):
        os.mkdir(osp.join(output_path, "texture_crops"))

    if not osp.exists(osp.join(output_path, "surface_texture_embeddings")):
        os.mkdir(osp.join(output_path, "surface_texture_embeddings"))

    # Computing RS embeddings on train set
    logging.info("Computing RS Embeddings on train set...")
    train_houses = load_houses_with_embeddings(conf, "train", "0.0", train_embeddings_path)
    rs_embeddings = compute_rs_mean_embeddings(conf, train_houses)

    # Load checkpoint
    predictor = TextureGenPredictor(
        conf=load_conf_eval(config_path=conf.texture_gen.texture_synth_conf),
        rgb_median_emb=conf.texture_gen.rgb_median_emb)
    predictor.load_checkpoint(checkpoint_path=conf.texture_gen.checkpoint_path)

    # Load houses
    house_keys = conf.get_data_list(split)
    houses = load_houses_with_embeddings(conf, split, conf.drop_fraction, osp.join(observed_surfaces_predictions_path, "surface_texture_embeddings"))

    # Process houses
    for i, (house_key, house) in enumerate(houses.items()):
        assert isinstance(house, House)
        logging.info("[%d/%d] Processing %s" % (i, len(houses), house_key))
        process_house(conf, house, predictor, rs_embeddings)
        save_house_crops(house, osp.join(output_path, "texture_crops", house_key))
        save_house_texture_embeddings(house, osp.join(output_path, "surface_texture_embeddings", house_key + ".json"))


""" BASIC USAGE OF THE CLASS PORTFOLIO """
# Change main directory to the main folder and import folders
### THIS CLASS DOES NOT DO MUCH, A QUICK WAY TO AFFECT ALL THE TIMESERIES
## OF THE SYMBOL IN A CERTAIN MANNER, like loading, or interval.
## It also contains info about the Symbol Itself that might be relevant.
import os
os.chdir("../../")
import import_folders
# Classical Libraries
import datetime as dt
import matplotlib.pyplot as plt

# Own graphical library
from graph_lib import gl
# Data Structures Data
import CSymbol as CSy
import CPortfolio as CPfl
# Import functions independent of DataStructure
import utilities_MQL5 as ul5

plt.close("all")
######## SELECT DATASET, SYMBOLS AND PERIODS ########
source = "MQL5" # Hanseatic  FxPro GCI Yahoo
[storage_folder, updates_folder] = ul5.get_foldersData(source = source)
################## Date info ###################
sdate_str = "01-01-2010"; edate_str = "21-12-2016"
sdate = dt.datetime.strptime(sdate_str, "%d-%m-%Y")
edate = dt.datetime.strptime(edate_str, "%d-%m-%Y")

#### Load the info about the available symbols #####
#Symbol_info = CSy.load_symbols_info(info_folder)
####### SELECT SYMBOLS AND PERIODS #################
symbolIDs =  ["EURUSD","EURGBP","USDCAD"]
periods = [1440]  # 1440  43200
period1 = periods[0]

####### LOAD SYMBOLS AND SET Properties   ###################
# Set the symbols and periods to load
myPortfolio = CPfl.Portfolio("BEST_PF",symbolIDs, periods)
# Download if needed.
#myPortfolio.update_symbols_csv_yahoo(sdate_str,edate_str,storage_folder)    # Load the symbols and periods
myPortfolio.set_csv(storage_folder)    # Load the symbols and periods

mySymbol_1 = myPortfolio.get_symbols([symbolIDs[0]])[0] # Obtain one of the symbols
myPortfolio.del_symbols([symbolIDs[2]])

# Create an independent symbol and add it to the portfolio
mySymbol_2 = CSy.CSymbol(symbolIDs[2],periods)
mySymbol_2.set_csv(storage_folder)
myPortfolio.add_symbols([mySymbol_2])

PF_symbolIDs = myPortfolio.get_symbolIDs() # ['DANSKE.CO', 'GE', 'HPQ']

## SET THINGS FOR ALL OF THEM
myPortfolio.set_interval(sdate,edate)
myPortfolio.set_seriesNames(["Close"])

########################## GETTING SINGALS ##########
# Obtaining some timeSeries
dates_all = myPortfolio.get_dates(period = periods[0])
close_some = myPortfolio.get_timeSeries(symbolIDs = symbolIDs)
return_open_all = myPortfolio.get_timeSeriesReturn(period = periods[0], seriesNames = ["Open"])

# Obtaining some indicators
EMA = myPortfolio.EMA(period = periods[0], n = 13)
ATR = myPortfolio.ATR(symbolIDs = symbolIDs ,n = 10)
MACD = myPortfolio.MACD()

######################## FILLING DATA  ##########################
#print "Filling Data"
#Cartera.fill_data()
#print "Data Filled"
###### WE CAN OPERATE DIRECTLY ON SYMBOLS #################

basic_joint_info = 0 # Obtaining info for the same period
trading_platform_several = 1 # Plot the trading platfom for all of them

###### FUNCTIONS TO OPERATE ON THE SAME PERIOD #############

if (trading_platform_several):
    # Plot for the 3 quitites, their Barchart and returns

    symbolIDs_pf = myPortfolio.get_symbolIDs()

#    dataTransform = ["intraday", opentime, closetime]
    folder_images = "../pics/gl/"
    gl.set_subplots(3,2)
     # TODO: Be able to automatize the shareX thing
    axeshare = None

    period = periods[0]
    for i in range(len(symbolIDs_pf)):
        symbolID = symbolIDs_pf[i]
        myTimeData = myPortfolio.get_symbols([symbolID])[0].get_timeData(period)
        returns = myTimeData.get_timeSeriesReturn(["Close"])
        dates = myTimeData.get_dates()
        AxesStyle = " - No xaxis"
        if (i == len(symbolIDs_pf) -1):
            AxesStyle = ""
        if (i == 0):
            title = "Bar Chart. " + str(symbolIDs) + r" . Price ($\$$)"
            title2 = "Return"
        else:
            title = ""
            title2 = ""
        ylabel =   symbolID + " (" + ul.period_dic[myTimeData.period] + ")"
        ax = gl.tradingBarChart(myTimeData,  legend = ["Close price"], color = "k",
                    nf = 1, sharex = axeshare, labels = [title,"",ylabel],
                    AxesStyle = "Normal" + AxesStyle)
#                    dataTransform = dataTransform)
        ax = gl.stem(dates, returns,  legend = ["Return"], color = "k",
                    nf = 1, sharex = axeshare, labels = [title2,"",""],
                    AxesStyle = "Normal" + AxesStyle + " - No yaxis")
#                    dataTransform = dataTransform)

        axeshare = ax

    gl.subplots_adjust(left=.09, bottom=.10, right=.90, top=.95, wspace=.10, hspace=0)
    image_name = "differentSymbols.png"
    gl.savefig(folder_images + image_name,
               dpi = 100, sizeInches = [30, 12])

Examples/1.1 Trading Information Classes/3.1 main_CPortfolio.py


#!/usr/bin/python3.7
# -*- coding: UTF-8 -*-
# author：Lucien
# 基础包: Log日志封装

import logging
from logging import handlers

class Logs(object):
    level_mapping = {
        'debug': logging.DEBUG,
        'info': logging.INFO,
        'warning': logging.WARNING,
        'error': logging.ERROR,
        'critical': logging.CRITICAL
    }  # 日志映射

    def __init__(self, filename, level='info', when='D', backCount=3,
                 fmt='%(asctime)s - [line:%(lineno)d] - %(levelname)s: %(message)s'):
        self.logger = logging.getLogger(filename)
        format_str = logging.Formatter(fmt)  # 设置格式
        self.logger.setLevel(self.level_mapping.get(level))  # 设置级别
        self.logger.propagate = False  # 关闭logger向上级传输
        stream = logging.StreamHandler()  # 流形式向屏幕输出
        stream.setFormatter(format_str)  # 流的显示的格式
        file = handlers.TimedRotatingFileHandler(filename=filename, when=when, backupCount=backCount,
                                                 encoding='utf-8')  # 往文件里写入
        # Calculate the real rollover interval, which is just the number of
        # seconds between rollovers.  Also set the filename suffix used when
        # a rollover occurs.  Current 'when' events supported:
        # S 秒
        # M 分
        # H 小时
        # D 天
        # W 每星期（interval==0时代表星期一）
        # midnight 凌晨
        file.setFormatter(format_str)  # 设置文件里写入的格式
        self.logger.addHandler(stream)  # 把对象加到logger里
        self.logger.addHandler(file)  # 把对象加到logger里

    def Logger(self):
        return self.logger

if __name__ == '__main__':
    log = Logs('logs').Logger()
    log.info('hello')


def quicksort(l):
    if(len(l) <= 1):
        return l
    less = []
    more = []
    n = len(l)
    pivot = 0
    for i in range(1, n):
        if(l[i] < l[pivot]):
            less.append(l[i])
        else:
            more.append(l[i])

    less = quicksort(less)
    more = quicksort(more)
    less.append(l[pivot])
    less.extend(more)
    return less

def mergesort(l):
    if(len(l) == 0 or len(l) == 1):
        return l

    left = mergesort(l[:int(len(l)/2)])
    right = mergesort(l[int(len(l)/2):])

    n = len(left)
    m = len(right)
    i = 0
    j = 0
    ans = []
    while(i < n and j < m):
        if(left[i] == right[j]):
            ans.append(left[i])
            ans.append(right[j])
            i = i+1
            j = j+1
        elif(left[i] < right[j]):
            ans.append(left[i])
            i = i+1
        else:
            ans.append(right[j])
            j = j+1

    while(i < n):
        ans.append(left[i])
        i = i+1
    while(j < m):
        ans.append(right[j])
        j = j+1

    return ans

def bubblesort(l):
    n = len(l)
    for i in range(n-1):
        for j in range(n-1-i):
            if(l[j] > l[j+1]):
                l[j], l[j+1] = l[j+1], l[j]
    return l

def comparativegraph():
    import random
    import time
    from matplotlib import pyplot as plt

    #bubtime = []
    mergetime = []
    quicktime = []
    for i in range(1, 5000):
        l = random.sample(range(100000), i)

        #start = time.time()
        #bubblesort(l.copy())
        #end = time.time()
        #bubtime.append(end-start)

        start = time.time()
        mergesort(l.copy())
        end = time.time()
        mergetime.append(end-start)

        start = time.time()
        quicksort(l.copy())
        end = time.time()
        quicktime.append(end-start)

    x = [i for i in range(1, 5000)]
    #plt.plot(x, bubtime)
    plt.plot(x, mergetime)
    plt.plot(x, quicktime)
    plt.legend(['merge', 'quick'])


import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

from utils.interpolate import interpolate

def local_deform(img, rm=None, dscale=4, a=2.0):

    if rm is None:
        map_size = (2* img.size(0), img.size(2) // dscale, img.size(3) // dscale)
        cm = 2.0 * (torch.rand(2 * img.size(0), 35, 3).to(img.device) - .5)
        rm = interpolate(cm.data, map_size).view(-1, 2, map_size[1], map_size[2]).detach()

    N = img.size(0)
    h, w = img.size(2), img.size(3)

    ig = torch.stack([
        torch.linspace(-1, 1, w).view(1, -1).repeat(h, 1),
        torch.linspace(-1, 1, h).view(-1, 1).repeat(1, w),
    ], 0).to(img.device)
    ig = ig.view(1, 2, h, w).repeat(N, 1, 1, 1)

    nrm = a * rm
    nrm[:, 0] -= torch.stack([l.mean() for l in nrm[:, 0]]).detach().view(-1, 1, 1)
    nrm[:, 1] -= torch.stack([l.mean() for l in nrm[:, 1]]).detach().view(-1, 1, 1)

    if (img.size(2) == rm.size(2)) and (img.size(3) == rm.size(3)):
        nig = ig.permute(0, 2, 3, 1).detach() + nrm.permute(0, 2, 3, 1)
    else:
        nig = ig.permute(0, 2, 3, 1).detach() + 1.0 * F.interpolate(nrm, (h, w), mode='bilinear').permute(0, 2, 3, 1)

    nimg = F.grid_sample(img, nig, mode='bilinear', padding_mode='border', align_corners=None)

    return nimg

def morphological(img, rm=None, dscale=4, a=10.0):

    if rm is None:
        map_size = (img.size(0), img.size(2) // dscale, img.size(3) // dscale)
        cm = 2.0 * (torch.rand(img.size(0), 35, 3).to(img.device) - .5)
        rm = interpolate(cm.data, map_size).view(-1, 1, map_size[1], map_size[2]).detach()

    h, w = img.size(2), img.size(3)

    dimg3 = F.max_pool2d(img, 3, 1, 1)
    dimg5 = F.max_pool2d(img, 5, 1, 2)
    eimg3 = -F.max_pool2d(-img, 3, 1, 1)
    eimg5 = -F.max_pool2d(-img, 5, 1, 2)
    simg = torch.cat([eimg5, eimg3, img, dimg3, dimg5], 1)

    nm = simg.size(1)

    centers = [float(i) / nm + 1.0 / (2 * nm) for i in range(nm)]  # [1.0 / 6, 1.0 / 3 + 1.0 / 6, 2.0 / 3 + 1.0 / 6]
    map1 = rm
    map1 = (map1 - map1.min()) / (map1.max() - map1.min())

    map3 = torch.cat([-(map1 - c) ** 2 for c in centers], 1)
    morph_mask = F.interpolate(map3, (h, w), mode='nearest')

    timg = simg * F.softmax(a * morph_mask, 1)
    timg = timg.sum(1).unsqueeze(1)

    return timg

def uncertainty_reduction(net, img, dscale=16, a=0.1, lr=1e-4, N=5, nc=50):

    sh, sw = img.size(2) // dscale, img.size(3) // dscale

    cm = 2.0 * (torch.rand(2 * img.size(0), nc, 4).to(img.device) - .5)
    cm[:, :, 2] *= .1
    cm[:, :, 3] += 10

    map_size = (2 * img.size(0), sh, sw)
    rm = interpolate(cm, map_size).view(-1, 2, map_size[1], map_size[2])
    rm = nn.Parameter(rm)

    toptimizer = torch.optim.Adam([rm], lr)

    dh = torch.Tensor([1, -1]).view(1, 1, 2, 1).repeat(2, 1, 1, 1).to(img.device)
    dw = dh.view(-1, 1, 1, 2)
    for i in range(N):
        toptimizer.zero_grad()

        nrm = a * rm #interpolate(cm, map_size).view(-1, 2, map_size[1], map_size[2])
        nimg = local_deform(img, nrm)

        output = net(nimg)

        eloss = F.conv2d(rm, dh, groups=2).norm() + F.conv2d(rm, dw, groups=2).norm()
        loss_val = -1.0 * (output.softmax(2)[:, :, :-1] ** 2).mean() + 0.0 * eloss

        loss_val.backward()
        toptimizer.step()

    nrm = a * rm #interpolate(cm, map_size).view(-1, 2, map_size[1], map_size[2])
    fimg = local_deform(img, nrm)

    return fimg


from contextlib import ExitStack
import functools
import importlib
import inspect
from itertools import chain
import os
import shlex
import shutil
import sys
import tempfile
import traceback
from typing import Optional, Sequence, Iterable, List, Type, Any, Callable
from pathlib import Path
from subprocess import check_call, run, PIPE, CompletedProcess, Popen

import click

@functools.lru_cache()
def mypy_cmd() -> Optional[Sequence[str]]:
    try:
        # preferably, use mypy from current python env
        import mypy  # noqa: F401 fine not to use it
    except ImportError:
        pass
    else:
        return [sys.executable, '-m', 'mypy']
    # ok, not ideal but try from PATH
    if shutil.which('mypy'):
        return ['mypy']
    warning("mypy not found, so can't check config with it. See https://github.com/python/mypy#readme if you want to install it and retry")
    return None

def run_mypy(cfg_path: Path) -> Optional[CompletedProcess]:
    # todo dunno maybe use the same mypy config in repository?
    # I'd need to install mypy.ini then??
    env = {**os.environ}
    mpath = env.get('MYPYPATH')
    mpath = str(cfg_path) + ('' if mpath is None else f':{mpath}')
    env['MYPYPATH'] = mpath

    cmd = mypy_cmd()
    if cmd is None:
        return None
    mres = run([
        *cmd,
        '--namespace-packages',
        '--color-output', # not sure if works??
        '--pretty',
        '--show-error-codes',
        '--show-error-context',
        '--check-untyped-defs',
        '-p', 'my.config',
    ], stderr=PIPE, stdout=PIPE, env=env)
    return mres

# use click.echo over print since it handles handles possible Unicode errors,
# strips colors if the output is a file
# https://click.palletsprojects.com/en/7.x/quickstart/#echoing
def eprint(x: str) -> None:
    # err=True prints to stderr
    click.echo(x, err=True)

def indent(x: str) -> str:
    return ''.join('   ' + l for l in x.splitlines(keepends=True))

OK  = '✅'
OFF = '🔲'

def info(x: str) -> None:
    eprint(OK + ' ' + x)

def error(x: str) -> None:
    eprint('❌ ' + x)

def warning(x: str) -> None:
    eprint('❗ ' + x) # todo yellow?

def tb(e: Exception) -> None:
    tb = ''.join(traceback.format_exception(Exception, e, e.__traceback__))
    sys.stderr.write(indent(tb))

def config_create() -> None:
    from .preinit import get_mycfg_dir
    mycfg_dir = get_mycfg_dir()

    created = False
    if not mycfg_dir.exists():
        # todo not sure about the layout... should I use my/config.py instead?
        my_config = mycfg_dir / 'my' / 'config' / '__init__.py'

        my_config.parent.mkdir(parents=True)
        my_config.write_text('''
### HPI personal config
## see
# https://github.com/karlicoss/HPI/blob/master/doc/SETUP.org#setting-up-modules
# https://github.com/karlicoss/HPI/blob/master/doc/MODULES.org
## for some help on writing your own config

# to quickly check your config, run:
# hpi config check

# to quickly check a specific module setup, run hpi doctor <module>, e.g.:
# hpi doctor my.reddit.rexport

### useful default imports
from my.core import Paths, PathIsh, get_files
###

# most of your configs will look like this:
class example:
    export_path: Paths = '/home/user/data/example_data_dir/'

### you can insert your own configuration below
### but feel free to delete the stuff above if you don't need ti
'''.lstrip())
        info(f'created empty config: {my_config}')
        created = True
    else:
        error(f"config directory '{mycfg_dir}' already exists, skipping creation")

    check_passed = config_ok()
    if not created or not check_passed:
        sys.exit(1)

# todo return the config as a result?
def config_ok() -> bool:
    errors: List[Exception] = []

    # at this point 'my' should already be imported, so doesn't hurt to extract paths from it
    import my
    try:
        paths: List[str] = list(my.__path__)
    except Exception as e:
        errors.append(e)
        error('failed to determine module import path')
        tb(e)
    else:
        info(f'import order: {paths}')

    # first try doing as much as possible without actually importing my.config
    from .preinit import get_mycfg_dir
    cfg_path = get_mycfg_dir()
    # alternative is importing my.config and then getting cfg_path from its __file__/__path__
    # not sure which is better tbh

    ## check we're not using stub config
    import my.core
    try:
        core_pkg_path = str(Path(my.core.__path__[0]).parent)
        if str(cfg_path).startswith(core_pkg_path):
            error(f'''
Seems that the stub config is used ({cfg_path}). This is likely not going to work.
See https://github.com/karlicoss/HPI/blob/master/doc/SETUP.org#setting-up-modules for more information
'''.strip())
            errors.append(RuntimeError('bad config path'))
    except Exception as e:
        errors.append(e)
        tb(e)
    else:
        info(f"config path : {cfg_path}")
    ##

    ## check syntax
    with tempfile.TemporaryDirectory() as td:
        # use a temporary directory, useful because
        # - compileall ignores -B, so always craps with .pyc files (annoyng on RO filesystems)
        # - compileall isn't following symlinks, just silently ignores them
        # note: ugh, annoying that copytree requires a non-existing dir before 3.8.
        # once we have min version 3.8, can use dirs_exist_ok=True param
        tdir = Path(td) / 'cfg'
        # this will resolve symlinks when copying
        shutil.copytree(cfg_path, tdir)
        # NOTE: compileall still returns code 0 if the path doesn't exist..
        # but in our case hopefully it's not an issue
        cmd = [sys.executable, '-m', 'compileall', '-q', str(tdir)]

        try:
            check_call(cmd)
            info('syntax check: ' + ' '.join(cmd))
        except Exception as e:
            errors.append(e)
            tb(e)
    ##

    ## check types
    mypy_res = run_mypy(cfg_path)
    if mypy_res is not None: # has mypy
        rc = mypy_res.returncode
        if rc == 0:
            info('mypy check  : success')
        else:
            error('mypy check: failed')
            errors.append(RuntimeError('mypy failed'))
            sys.stderr.write(indent(mypy_res.stderr.decode('utf8')))
        sys.stderr.write(indent(mypy_res.stdout.decode('utf8')))
    ##

    ## finally, try actually importing the config (it should use same cfg_path)
    try:
        import my.config
    except Exception as e:
        errors.append(e)
        error("failed to import the config")
        tb(e)
    ##

    if len(errors) > 0:
        error(f'config check: {len(errors)} errors')
        return False
    else:
        # note: shouldn't exit here, might run something else
        info('config check: success!')
        return True

from .util import HPIModule, modules
def _modules(*, all: bool=False) -> Iterable[HPIModule]:
    skipped = []
    for m in modules():
        if not all and m.skip_reason is not None:
            skipped.append(m.name)
        else:
            yield m
    if len(skipped) > 0:
        warning(f'Skipped {len(skipped)} modules: {skipped}. Pass --all if you want to see them.')

def modules_check(*, verbose: bool, list_all: bool, quick: bool, for_modules: List[str]) -> None:
    if len(for_modules) > 0:
        # if you're checking specific modules, show errors
        # hopefully makes sense?
        verbose = True
    vw = '' if verbose else '; pass --verbose to print more information'

    tabulate_warnings()

    import contextlib

    from .common import quick_stats
    from .util import get_stats, HPIModule
    from .stats import guess_stats
    from .error import warn_my_config_import_error

    mods: Iterable[HPIModule]
    if len(for_modules) == 0:
        mods = _modules(all=list_all)
    else:
        mods = [HPIModule(name=m, skip_reason=None) for m in for_modules]

    # todo add a --all argument to disregard is_active check?
    for mr in mods:
        skip = mr.skip_reason
        m    = mr.name
        if skip is not None:
            eprint(f'{OFF} {click.style("SKIP", fg="yellow")}: {m:<50} {skip}')
            continue

        try:
            mod = importlib.import_module(m)  # noqa: F841
        except Exception as e:
            # todo more specific command?
            error(f'{click.style("FAIL", fg="red")}: {m:<50} loading failed{vw}')
            # check that this is an import error in particular, not because
            # of a ModuleNotFoundError because some dependency wasn't installed
            if isinstance(e, (ImportError, AttributeError)):
                warn_my_config_import_error(e)
            if verbose:
                tb(e)
            continue

        info(f'{click.style("OK", fg="green")}  : {m:<50}')
        # first try explicitly defined stats function:
        stats = get_stats(m)
        if stats is None:
            # then try guessing.. not sure if should log somehow?
            stats = guess_stats(m, quick=quick)

        if stats is None:
            eprint("       - no 'stats' function, can't check the data")
            # todo point to a readme on the module structure or something?
            continue

        quick_context = quick_stats() if quick else contextlib.nullcontext()

        try:
            kwargs = {}
            if callable(stats) and 'quick' in inspect.signature(stats).parameters:
                kwargs['quick'] = quick
            with quick_context:
                res = stats(**kwargs)
            assert res is not None, 'stats() returned None'
        except Exception as ee:
            warning(f'     - {click.style("stats:", fg="red")}                      computing failed{vw}')
            if verbose:
                tb(ee)
        else:
            info(f'    - stats: {res}')

def list_modules(*, list_all: bool) -> None:
    # todo add a --sort argument?
    tabulate_warnings()

    for mr in _modules(all=list_all):
        m    = mr.name
        sr   = mr.skip_reason
        if sr is None:
            pre = OK
            suf = ''
        else:
            pre = OFF
            suf = f' {click.style(f"[disabled: {sr}]", fg="yellow")}'

        click.echo(f'{pre} {m:50}{suf}')

def tabulate_warnings() -> None:
    '''
    Helper to avoid visual noise in hpi modules/doctor
    '''
    import warnings
    orig = warnings.formatwarning

    def override(*args, **kwargs) -> str:
        res = orig(*args, **kwargs)
        return ''.join('  ' + x for x in res.splitlines(keepends=True))
    warnings.formatwarning = override
    # TODO loggers as well?

def _requires(modules: Sequence[str]) -> Sequence[str]:
    from .discovery_pure import module_by_name
    mods = [module_by_name(module) for module in modules]
    res = []
    for mod in mods:
        if mod.legacy is not None:
            warning(mod.legacy)

        reqs = mod.requires
        if reqs is None:
            warning(f"Module {mod.name} has no REQUIRES specification")
            continue
        for r in reqs:
            if r not in res:
                res.append(r)
    return res

def module_requires(*, module: Sequence[str]) -> None:
    if isinstance(module, str):
        # legacy behavior, used to take a since argument
        module = [module]
    rs = [f"'{x}'" for x in _requires(modules=module)]
    eprint(f'dependencies of {module}')
    for x in rs:
        click.echo(x)

def module_install(*, user: bool, module: Sequence[str], parallel: bool=False, break_system_packages: bool=False) -> None:
    if isinstance(module, str):
        # legacy behavior, used to take a since argument
        module = [module]

    requirements = _requires(module)

    if len(requirements) == 0:
        warning('requirements list is empty, no need to install anything')
        return

    pre_cmd = [
        sys.executable, '-m', 'pip',
        'install',
        *(['--user'] if user else []), # todo maybe instead, forward all the remaining args to pip?
        *(['--break-system-packages'] if break_system_packages else []), # https://peps.python.org/pep-0668/
    ]

    cmds = []
    # disable parallel on windows, sometimes throws a
    # '[WinError 32] The process cannot access the file because it is being used by another process'
    # same on mac it seems? possible race conditions which are hard to debug?
    # WARNING: Error parsing requirements for sqlalchemy: [Errno 2] No such file or directory: '/Users/runner/work/HPI/HPI/.tox/mypy-misc/lib/python3.7/site-packages/SQLAlchemy-2.0.4.dist-info/METADATA'
    if parallel and sys.platform not in ['win32', 'cygwin', 'darwin']:
        # todo not really sure if it's safe to install in parallel like this
        # but definitely doesn't hurt to experiment for e.g. mypy pipelines
        # pip has '--use-feature=fast-deps', but it doesn't really work
        # I think it only helps for pypi artifacts (not git!),
        # and only if they weren't cached
        for r in requirements:
            cmds.append(pre_cmd + [r])
    else:
        if parallel:
            warning('parallel install is not supported on this platform, installing sequentially...')
        # install everything in one cmd
        cmds.append(pre_cmd + list(requirements))

    with ExitStack() as exit_stack:
        popens = []
        for cmd in cmds:
            eprint('Running: ' + ' '.join(map(shlex.quote, cmd)))
            popen = exit_stack.enter_context(Popen(cmd))
            popens.append(popen)

        for popen in popens:
            ret = popen.wait()
            assert ret == 0, popen

def _ui_getchar_pick(choices: Sequence[str], prompt: str = 'Select from: ') -> int:
    '''
    Basic menu allowing the user to select one of the choices
    returns the index the user chose
    '''
    assert len(choices) > 0, 'Didnt receive any choices to prompt!'
    eprint(prompt + '\n')

    # prompts like 1,2,3,4,5,6,7,8,9,a,b,c,d,e,f...
    chr_offset = ord('a') - 10

    # dict from key user can press -> resulting index
    result_map = {}
    for i, opt in enumerate(choices, 1):
        char: str = str(i) if i < 10 else chr(i + chr_offset)
        result_map[char] = i - 1
        eprint(f'\t{char}. {opt}')

    eprint('')
    while True:
        ch = click.getchar()
        if ch not in result_map:
            eprint(f'{ch} not in {list(result_map.keys())}')
            continue
        return result_map[ch]

def _locate_functions_or_prompt(qualified_names: List[str], prompt: bool = True) -> Iterable[Callable[..., Any]]:
    from .query import locate_qualified_function, QueryException
    from .stats import is_data_provider

    # if not connected to a terminal, can't prompt
    if not sys.stdout.isatty():
        prompt = False

    for qualname in qualified_names:
        try:
            # common-case
            yield locate_qualified_function(qualname)
        except QueryException as qr_err:
            # maybe the user specified a module name instead of a function name?
            # try importing the name the user specified as a module and prompt the
            # user to select a 'data provider' like function
            try:
                mod = importlib.import_module(qualname)
            except Exception:
                eprint(f"During fallback, importing '{qualname}' as module failed")
                raise qr_err

            # find data providers in this module
            data_providers = [f for _, f in inspect.getmembers(mod, inspect.isfunction) if is_data_provider(f)]
            if len(data_providers) == 0:
                eprint(f"During fallback, could not find any data providers in '{qualname}'")
                raise qr_err
            else:
                # was only one data provider-like function, use that
                if len(data_providers) == 1:
                    yield data_providers[0]
                else:
                    choices = [f.__name__ for f in data_providers]
                    if prompt is False:
                        # there's more than one possible data provider in this module,
                        # STDOUT is not a TTY, can't prompt
                        eprint("During fallback, more than one possible data provider, can't prompt since STDOUT is not a TTY")
                        eprint("Specify one of:")
                        for funcname in choices:
                            eprint(f"\t{qualname}.{funcname}")
                        raise qr_err
                    # prompt the user to pick the function to use
                    chosen_index = _ui_getchar_pick(choices, f"Which function should be used from '{qualname}'?")
                    # respond to the user, so they know something has been picked
                    eprint(f"Selected '{choices[chosen_index]}'")
                    yield data_providers[chosen_index]

def _warn_exceptions(exc: Exception) -> None:
    from my.core.common import LazyLogger
    logger = LazyLogger('CLI', level='warning')

    logger.exception(f'hpi query: {exc}')

# handle the 'hpi query' call
# can raise a QueryException, caught in the click command
def query_hpi_functions(
    *,
    output: str = 'json',
    stream: bool = False,
    qualified_names: List[str],
    order_key: Optional[str],
    order_by_value_type: Optional[Type],
    after: Any,
    before: Any,
    within: Any,
    reverse: bool = False,
    limit: Optional[int],
    drop_unsorted: bool,
    wrap_unsorted: bool,
    warn_exceptions: bool,
    raise_exceptions: bool,
    drop_exceptions: bool,
) -> None:
    from .query_range import select_range, RangeTuple
    import my.core.error as err

    # chain list of functions from user, in the order they wrote them on the CLI
    input_src = chain(*(f() for f in _locate_functions_or_prompt(qualified_names)))

    res = select_range(
        input_src,
        order_key=order_key,
        order_by_value_type=order_by_value_type,
        unparsed_range=RangeTuple(after=after, before=before, within=within),
        reverse=reverse,
        limit=limit,
        drop_unsorted=drop_unsorted,
        wrap_unsorted=wrap_unsorted,
        warn_exceptions=warn_exceptions,
        warn_func=_warn_exceptions,
        raise_exceptions=raise_exceptions,
        drop_exceptions=drop_exceptions)

    if output == 'json':
        from .serialize import dumps

        if stream:
            for item in res:
                # use sys.stdout directly
                # the overhead form click.echo isn't a *lot*, but when called in a loop
                # with potentially millions of items it makes a noticeable difference
                sys.stdout.write(dumps(item))
                sys.stdout.write('\n')
            sys.stdout.flush()
        else:
            click.echo(dumps(list(res)))
    elif output == 'pprint':
        from pprint import pprint

        if stream:
            for item in res:
                pprint(item)
        else:
            pprint(list(res))
    elif output == 'gpx':
        from my.location.common import locations_to_gpx

        # if user didn't specify to ignore exceptions, warn if locations_to_gpx
        # cannot process the output of the command. This can be silenced by
        # passing --drop-exceptions
        if not raise_exceptions and not drop_exceptions:
            warn_exceptions = True

        # can ignore the mypy warning here, locations_to_gpx yields any errors
        # if you didnt pass it something that matches the LocationProtocol
        for exc in locations_to_gpx(res, sys.stdout):   # type: ignore[arg-type]
            if warn_exceptions:
                _warn_exceptions(exc)
            elif raise_exceptions:
                raise exc
            elif drop_exceptions:
                pass
        sys.stdout.flush()
    else:
        res = list(res)  # type: ignore[assignment]
        # output == 'repl'
        eprint(f"\nInteract with the results by using the {click.style('res', fg='green')} variable\n")
        try:
            import IPython  # type: ignore[import]
        except ModuleNotFoundError:
            eprint("'repl' typically uses ipython, install it with 'python3 -m pip install ipython'. falling back to stdlib...")
            import code
            code.interact(local=locals())
        else:
            IPython.embed()

@click.group()
@click.option("--debug", is_flag=True, default=False, help="Show debug logs")
def main(debug: bool) -> None:
    '''
    Human Programming Interface

    Tool for HPI
    Work in progress, will be used for config management, troubleshooting & introspection
    '''
    # should overwrite anything else in LOGGING_LEVEL_HPI
    if debug:
        os.environ['LOGGING_LEVEL_HPI'] = 'debug'

    # for potential future reference, if shared state needs to be added to groups
    # https://click.palletsprojects.com/en/7.x/commands/#group-invocation-without-command
    # https://click.palletsprojects.com/en/7.x/commands/#multi-command-chaining

    # acts as a contextmanager of sorts - any subcommand will then run
    # in something like /tmp/hpi_temp_dir
    # to avoid importing relative modules by accident during development
    # maybe can be removed later if there's more test coverage/confidence that nothing
    # would happen?

    # use a particular directory instead of a random one, since
    # click being decorator based means its more complicated
    # to run things at the end (would need to use a callback or pass context)
    # https://click.palletsprojects.com/en/7.x/commands/#nested-handling-and-contexts

    tdir: str = os.path.join(tempfile.gettempdir(), 'hpi_temp_dir')
    if not os.path.exists(tdir):
        os.makedirs(tdir)
    os.chdir(tdir)

@functools.lru_cache(maxsize=1)
def _all_mod_names() -> List[str]:
    """Should include all modules, in case user is trying to diagnose issues"""
    # sort this, so that the order doesn't change while tabbing through
    return sorted([m.name for m in modules()])

def _module_autocomplete(ctx: click.Context, args: Sequence[str], incomplete: str) -> List[str]:
    return [m for m in _all_mod_names() if m.startswith(incomplete)]

@main.command(name='doctor', short_help='run various checks')
@click.option('--verbose/--quiet', default=False, help='Print more diagnostic information')
@click.option('--all', 'list_all', is_flag=True, help='List all modules, including disabled')
@click.option('-q', '--quick', is_flag=True, help='Only run partial checks (first 100 items)')
@click.option('-S', '--skip-config-check', 'skip_conf', is_flag=True, help='Skip configuration check')
@click.argument('MODULE', nargs=-1, required=False, shell_complete=_module_autocomplete)
def doctor_cmd(verbose: bool, list_all: bool, quick: bool, skip_conf: bool, module: Sequence[str]) -> None:
    '''
    Run various checks

    MODULE is one or more specific module names to check (e.g. my.reddit.rexport)
    Otherwise, checks all modules
    '''
    if not skip_conf:
        config_ok()
    # TODO check that it finds private modules too?
    modules_check(verbose=verbose, list_all=list_all, quick=quick, for_modules=list(module))

@main.group(name='config', short_help='work with configuration')
def config_grp() -> None:
    '''Act on your HPI configuration'''
    pass

@config_grp.command(name='check', short_help='check config')
def config_check_cmd() -> None:
    '''Check your HPI configuration file'''
    ok = config_ok()
    sys.exit(0 if ok else False)

@config_grp.command(name='create', short_help='create user config')
def config_create_cmd() -> None:
    '''Create user configuration file for HPI'''
    config_create()

@main.command(name='modules', short_help='list available modules')
@click.option('--all', 'list_all', is_flag=True, help='List all modules, including disabled')
def module_cmd(list_all: bool) -> None:
    '''List available modules'''
    list_modules(list_all=list_all)

@main.group(name='module', short_help='module management')
def module_grp() -> None:
    '''Module management'''
    pass

@module_grp.command(name='requires', short_help='print module reqs')
@click.argument('MODULES', shell_complete=_module_autocomplete, nargs=-1, required=True)
def module_requires_cmd(modules: Sequence[str]) -> None:
    '''
    Print MODULES requirements

    MODULES is one or more specific module names (e.g. my.reddit.rexport)
    '''
    module_requires(module=modules)

@module_grp.command(name='install', short_help='install module deps')
@click.option('--user', is_flag=True, help='same as pip --user')
@click.option('--parallel', is_flag=True, help='EXPERIMENTAL. Install dependencies in parallel.')
@click.option('-B',
              '--break-system-packages',
              is_flag=True,
              help='Bypass PEP 668 and install dependencies into the system-wide python package directory.')
@click.argument('MODULES', shell_complete=_module_autocomplete, nargs=-1, required=True)
def module_install_cmd(user: bool, parallel: bool, break_system_packages: bool, modules: Sequence[str]) -> None:
    '''
    Install dependencies for modules using pip

    MODULES is one or more specific module names (e.g. my.reddit.rexport)
    '''
    # todo could add functions to check specific module etc..
    module_install(user=user, module=modules, parallel=parallel, break_system_packages=break_system_packages)

@main.command(name='query', short_help='query the results of a HPI function')
@click.option('-o',
              '--output',
              default='json',
              type=click.Choice(['json', 'pprint', 'repl', 'gpx']),
              help='what to do with the result [default: json]')
@click.option('-s',
              '--stream',
              default=False,
              is_flag=True,
              help='stream objects from the data source instead of printing a list at the end')
@click.option('-k',
              '--order-key',
              default=None,
              type=str,
              help='order by an object attribute or dict key on the individual objects returned by the HPI function')
@click.option('-t',
              '--order-type',
              default=None,
              type=click.Choice(['datetime', 'date', 'int', 'float']),
              help='order by searching for some type on the iterable')
@click.option('-a',
              '--after',
              default=None,
              type=str,
              help='while ordering, filter items for the key or type larger than or equal to this')
@click.option('-b',
              '--before',
              default=None,
              type=str,
              help='while ordering, filter items for the key or type smaller than this')
@click.option('-w',
              '--within',
              default=None,
              type=str,
              help="a range 'after' or 'before' to filter items by. see above for further explanation")
@click.option('-r',
              '--recent',
              default=None,
              type=str,
              help="a shorthand for '--order-type datetime --reverse --before now --within'. e.g. --recent 5d")
@click.option('--reverse/--no-reverse',
              default=False,
              help='reverse the results returned from the functions')
@click.option('-l',
              '--limit',
              default=None,
              type=int,
              help='limit the number of items returned from the (functions)')
@click.option('--drop-unsorted',
              default=False,
              is_flag=True,
              help="if the order of an item can't be determined while ordering, drop those items from the results")
@click.option('--wrap-unsorted',
              default=False,
              is_flag=True,
              help="if the order of an item can't be determined while ordering, wrap them into an 'Unsortable' object")
@click.option('--warn-exceptions',
              default=False,
              is_flag=True,
              help="if any errors are returned, print them as errors on STDERR")
@click.option('--raise-exceptions',
              default=False,
              is_flag=True,
              help="if any errors are returned (as objects, not raised) from the functions, raise them")
@click.option('--drop-exceptions',
              default=False,
              is_flag=True,
              help='ignore any errors returned as objects from the functions')
@click.argument('FUNCTION_NAME', nargs=-1, required=True, shell_complete=_module_autocomplete)
def query_cmd(
    function_name: Sequence[str],
    output: str,
    stream: bool,
    order_key: Optional[str],
    order_type: Optional[str],
    after: Optional[str],
    before: Optional[str],
    within: Optional[str],
    recent: Optional[str],
    reverse: bool,
    limit: Optional[int],
    drop_unsorted: bool,
    wrap_unsorted: bool,
    warn_exceptions: bool,
    raise_exceptions: bool,
    drop_exceptions: bool,
) -> None:
    '''
    This allows you to query the results from one or more functions in HPI

    By default this runs with '-o json', converting the results
    to JSON and printing them to STDOUT

    You can specify '-o pprint' to just print the objects using their
    repr, or '-o repl' to drop into a ipython shell with access to the results

    While filtering using --order-key datetime, the --after, --before and --within
    flags parse the input to their datetime and timedelta equivalents. datetimes can
    be epoch time, the string 'now', or an date formatted in the ISO format. timedelta
    (durations) are parsed from a similar format to the GNU 'sleep' command, e.g.
    1w2d8h5m20s -> 1 week, 2 days, 8 hours, 5 minutes, 20 seconds

    As an example, to query reddit comments I've made in the last month

    \b
    hpi query --order-type datetime --before now --within 4w my.reddit.all.comments
    or...
    hpi query --recent 4w my.reddit.all.comments

    \b
    Can also query within a range. To filter comments between 2016 and 2018:
    hpi query --order-type datetime --after '2016-01-01' --before '2019-01-01' my.reddit.all.comments
    '''

    from datetime import datetime, date

    chosen_order_type: Optional[Type]
    if order_type == "datetime":
        chosen_order_type = datetime
    elif order_type == "date":
        chosen_order_type = date
    elif order_type == "int":
        chosen_order_type = int
    elif order_type == "float":
        chosen_order_type = float
    else:
        chosen_order_type = None

    if recent is not None:
        before = "now"
        chosen_order_type = chosen_order_type or datetime  # dont override if the user specified date
        within = recent
        reverse = not reverse

    from .query import QueryException

    try:
        query_hpi_functions(
            output=output,
            stream=stream,
            qualified_names=list(function_name),
            order_key=order_key,
            order_by_value_type=chosen_order_type,
            after=after,
            before=before,
            within=within,
            reverse=reverse,
            limit=limit,
            drop_unsorted=drop_unsorted,
            wrap_unsorted=wrap_unsorted,
            warn_exceptions=warn_exceptions,
            raise_exceptions=raise_exceptions,
            drop_exceptions=drop_exceptions)
    except QueryException as qe:
        eprint(str(qe))
        sys.exit(1)

# todo: add more tests?
# its standard click practice to have the function click calls be a separate
# function from the decorated function, as it allows the application-specific code to be
# more testable. also allows hpi commands to be imported and called manually from
# other python code

def test_requires() -> None:
    from click.testing import CliRunner
    result = CliRunner().invoke(main, ['module', 'requires', 'my.github.ghexport', 'my.browser.export'])
    assert result.exit_code == 0
    assert "github.com/karlicoss/ghexport" in result.output
    assert "browserexport" in result.output

if __name__ == '__main__':
    # prog_name is so that if this is invoked with python -m my.core
    # this still shows hpi in the help text
    main(prog_name='hpi')


from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig, StoppingCriteria, StoppingCriteriaList
import torch

test_data = load_dataset("json", data_files="/root/data/data/openai_humaneval/HumanEval.jsonl")
# test_data = load_dataset("json", data_files="/root/data/data/humaneval-x/data/python/data/humaneval.jsonl")

print(test_data)

# tok = AutoTokenizer.from_pretrained("/root/data/bloom", padding_side="left")
tok = AutoTokenizer.from_pretrained("/root/data/bloom-560m", padding_side="left")
# tok = AutoTokenizer.from_pretrained("/root/data/starchat", padding_side="left")
# tokenizer.pad_token = tokenizer.eos_token
tok.pad_token = tok.eos_token
print("init tok done")

# model = AutoModelForCausalLM.from_pretrained("/root/data/bloom").cuda()
model = AutoModelForCausalLM.from_pretrained("/root/data/bloom-560m").cuda()
# model = AutoModelForCausalLM.from_pretrained("/root/data/starchat").half().cuda()
model = model.eval()
print("init model done")

system_token = "<|system|>"
user_token = "<|user|>"
assistant_token = "<|assistant|>"
end_token = "<|end|>"
max_length = 2048
stop_token_ids = [tok.eos_token_id, tok(end_token)["input_ids"][0], tok(user_token)["input_ids"][0]]
class StopOnTokens(StoppingCriteria):
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
        for stop_id in stop_token_ids:
            if input_ids[0][-1] == stop_id:
                return True
        return False
stop = StopOnTokens()
generation_config = GenerationConfig(
        temperature=0.2,
        top_k=50,
        top_p=0.95,
        repetition_penalty=1.2,
        do_sample=True,
        pad_token_id=tok.eos_token_id,
        # eos_token_id=tok.convert_tokens_to_ids(dialogue_template.end_token),
        min_new_tokens=32,
        max_new_tokens=256,
        # stopping_criteria=StoppingCriteriaList([stop])
    )

def gen_code(model, tok, prompt):
    prompts = tok(
        prompt,
        max_length=max_length,
        padding="max_length",
        truncation=True,
        return_tensors="pt",
    )

    input_ids = torch.tensor(prompts['input_ids']).cuda()
    attention_mask =  torch.tensor(prompts['attention_mask']).cuda()

    output = model.generate(input_ids=input_ids, generation_config=generation_config)
    # output = model.generate(
    #     input_ids=input_ids,
    #     attention_mask=attention_mask,
    #     max_length=max_length,
    #     do_sample=True,
    #     top_k=50,
    #     top_p=0.95,
    #     num_beams=5,
    #     eos_token_id=tok.eos_token_id,
    #     num_return_sequences=1,
    #     early_stopping=True,
    #     length_penalty=2.5)

    return tok.decode(output[0], skip_special_tokens=True)

train_data = test_data['train']

with open("./test_problem_top6_sc_prompt.txt", 'w') as wp:
    for idx, data in enumerate(train_data):
        # print(data)
        # for d in data:
            # print("-----")
            # print(d)
            # print(data[d])

        system_msg = "Below are a series of dialogues between various people and an AI assistant.\
The AI tries to be helpful, polite, honest, sophisticated, emotionally aware, and humble-but-knowledgeable.\
The assistant is happy to help with almost anything, and will do its best to understand exactly what is needed.\
It also tries to avoid giving false or misleading information, and it caveats when it isn’t entirely sure about the right answer.\
That said, the assistant is practical and really does its best, and doesn’t let caution get too much in the way of being useful.\
-----\n"

        prompt = system_token + "\n" + system_msg + end_token + "\n" + user_token + "\n" + data['prompt'] + end_token + "\n" + assistant_token + "\n"
        # prompt = "continue write the code below\n\n \"\"\"" + data['prompt'] + "\"\"\""
        code = gen_code(model, tok, prompt)
        print("-----------")
        print("prompt:", prompt)
        # print("prompt:", data['prompt'])
        # print("[code]:", code)
        print("[code]:", code.replace(prompt, ""))
        print("ans: \n", data['canonical_solution'])
        print("example_test:", data['test'])
        wp.write("----------------\n")
        wp.write(code + "\n")
        if idx > 6:
            break
        # break

print(len(train_data))


import pygame, math, sys
from time import sleep
pygame.init()

X = 900
Y = 650

WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
GREY = (200, 200, 200)
ORANGE = (200, 100, 50)
RED = (255,0,0)
RED_2 = (237,41,57)
GREEN = (0,255,0)
TRANS = (1, 1, 1)

font = pygame.font.SysFont("Arial", 12)
font2 = pygame.font.SysFont("Arial BOLD", 24)
screen = pygame.display.set_mode((X, Y))
clock = pygame.time.Clock()
pop = 7

background2 = pygame.image.load('background2_forces.jpg')
background3 = pygame.image.load('background3_forces.jpg')

horiz_plane = pygame.image.load("horiz_plane.png")
horiz_plane = pygame.transform.scale(horiz_plane,(60,60))

class Slider():
    def __init__(self, name, maxi, mini, pos, unit):
        self.val = (maxi+mini)/2  # start value
        self.maxi = maxi  # maximum at slider position right
        self.mini = mini  # minimum at slider position left
        self.xpos = X/2-40  # x-location on screen
        self.x_max = self.xpos+80
        self.ypos = pos
        self.surf = pygame.surface.Surface((165, 50))
        self.hit = False  # the hit attribute indicates slider movement due to mouse interaction
        self.unit = unit

        self.txt_surf = font.render(name, 1, BLACK)
        self.txt_rect = self.txt_surf.get_rect(center=(50, 15))

        # Static graphics - slider background #
        pygame.draw.rect(self.surf, GREY, [0, 0, 165, 50], 3)
        pygame.draw.rect(self.surf, ORANGE, [10, 10, 80, 10], 0)
        pygame.draw.rect(self.surf, WHITE, [10, 30, 80, 5], 0)

        self.surf.blit(self.txt_surf, self.txt_rect)  # this surface never changes

        # dynamic graphics - button surface #
        self.button_surf = pygame.surface.Surface((20, 20))
        self.button_surf.fill(TRANS)
        self.button_surf.set_colorkey(TRANS)
        pygame.draw.circle(self.button_surf, BLACK, (10, 10), 6, 0)
        pygame.draw.circle(self.button_surf, ORANGE, (10, 10), 4, 0)

    def draw(self):
        # static
        surf = self.surf.copy()
        # dynamic
        pos = (10+int((self.val-self.mini)/(self.maxi-self.mini)*80), 33)
        self.button_rect = self.button_surf.get_rect(center=pos)
        surf.blit(self.button_surf, self.button_rect)
        self.button_rect.move_ip(self.xpos, self.ypos)
        # screen
        mouseX,mouseY=pygame.mouse.get_pos()
        screen.blit(surf, (self.xpos, self.ypos))
        value = font.render("= "+str(round(self.val,2))+self.unit,False,WHITE)
        screen.blit(value,(self.xpos+100,self.ypos+15))

    def move(self):
        self.val = (pygame.mouse.get_pos()[0] - self.xpos - 10) / 80 * (self.maxi - self.mini) + self.mini
        if self.val < self.mini:
            self.val = self.mini
        if self.val > self.maxi:
            self.val = self.maxi

def start_button_display():
    mousex, mousey = pygame.mouse.get_pos()
    if 10 <= mousex <= 10+50 and  75 <= mousey <= 75+25 and pygame.mouse.get_pressed():
        pygame.draw.rect(screen, GREEN, [10-pop/2, 75-pop/2, 50+pop, 25+pop])
        pygame.draw.rect(screen, WHITE, (10, 75, 50, 25))
        return True
    pygame.draw.rect(screen, GREEN, [10-pop/2, 75-pop/2, 50+pop, 25+pop])
    return False

def plane_startup(planeY,count,component):
    if count == 0: #BASE CASE
        return planeY

    if count % 2 ==1:
        planeY-=5
    else:
        planeY+=5
    count-=1
    screen.blit(horiz_plane,(X/2,Y-100+planeY))
    pygame.display.update()
    if component == 'Vertical':
        screen.blit(background2,(0,0))
    elif component == 'Horizontal':
        screen.blit(background3,(0,0))
    sleep(.05)
    return plane_startup(planeY,count,component)

def call_sliders(slides,component):
    DisplayList=[]
    planeX=X/2
    planeY=plane_startup(0,20, component)
    while True:
        if component == 'Vertical':
            force_gravity = slides[0].val*9.8
            a=(force_gravity - slides[1].val)/slides[0].val
            if start_button_display() and a>0:
                return_list=[]
                for i in range(len(slides)):
                    return_list.append(slides[i].val)
                return return_list
            elif a<=0:
                pygame.draw.rect(screen, RED, [10-pop/2, 75-pop/2, 50+pop, 25+pop])
                static_text = font2.render('Impossible Acceleration value ', False, WHITE)
                static_text2 = font.render('Adjust Air Resistance or Mass!', False, RED_2)
                screen.blit(static_text,(X/2-120+70-30, Y/2+140+30-5))
                screen.blit(static_text2,(X/2-120+50, Y/2+150+30))

        elif component == 'Horizontal':
            force_friction = slides[0].val*9.8*slides[2].val
            a=(slides[1].val - force_friction)/slides[0].val
            if start_button_display() and a>0:
                return_list=[]
                for i in range(len(slides)):
                    return_list.append(slides[i].val)
                return return_list
            elif a<=0:
                pygame.draw.rect(screen, RED, [10-pop/2, 75-pop/2, 50+pop, 25+pop])
                static_text = font2.render('Impossible Acceleration value ', False, WHITE)
                static_text2 = font.render('Adjust Mass, Coef or Applied Force!', False, RED_2)
                screen.blit(static_text,(X/2-120+70-30, Y/2+140+30-5))
                screen.blit(static_text2,(X/2-120+50, Y/2+150+30))

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            elif event.type == pygame.MOUSEBUTTONDOWN:
                pos = pygame.mouse.get_pos()
                for s in slides:
                    if s.button_rect.collidepoint(pos):
                        s.hit =  True
            elif event.type == pygame.MOUSEBUTTONUP:
                for s in slides:
                    s.hit = False

        for s in slides:
            if s.hit:
                s.move()

        pygame.display.update()

        if component == 'Vertical':
            screen.blit(background2,(0,0))
        elif component == 'Horizontal':
            screen.blit(background3,(0,0))
        planeX+=1
        if planeX-120>X: planeX=0-120
        screen.blit(horiz_plane,(planeX,Y-100))
        for s in slides:
            s.draw()


from flask import request, jsonify
from mysqlAccess import mysqlConnection, mysqlCloseConnection, mysqlQuery, mysqlInsert, mysqlUpdate
from configApp import app

@app.route('/participacoes', methods=['GET'])
def participacoes():
    accessPoints = ["all", "att", "describe", "insert", "remove"]
    return jsonify(accessPoints)

@app.route('/participacoes/all', methods=['GET'])
def participacoesAll():
    conn = mysqlConnection()
    query = "SELECT * FROM participacoes;"
    records = mysqlQuery(conn, query)
    mysqlCloseConnection(conn)
    return jsonify(records)

@app.route('/participacoes/att', methods=['GET'])
def participacoesAtt():
    conn = mysqlConnection()
    cursor = conn.cursor()
    cursor.execute(f'SELECT * FROM participacoes;')
    column_names = [i[0] for i in cursor.description]
    data = cursor.fetchall()
    rows = []
    for row in data:
        row_dict = {}
        for i in range(len(column_names)):
            row_dict[column_names[i]] = row[i]
        rows.append(row_dict)
    mysqlCloseConnection(conn)
    return jsonify(rows)

@app.route('/participacoes/describe', methods=['GET'])
def participacoesDescribe():
    conn = mysqlConnection()
    if 'idEvento' not in request.args and 'email' not in request.args:
        return "Error: No idEvento and email fields provided."
    elif 'idEvento' not in request.args:
        return "Error: No idEvento field provided. Please specify an id."
    elif 'email' not in request.args:
        return "Error: No email field provided. Please specify an email."

    idEvento = request.args['idEvento']
    email = request.args['email']
    query = "SELECT * FROM participacoes WHERE idEvento = '"+idEvento+"' AND email='"+email+"';"
    records = mysqlQuery(conn, query)
    mysqlCloseConnection(conn)
    return jsonify(records)

@app.route('/participacoes/insert', methods=['GET'])
def participacoesInsert():
    conn = mysqlConnection()
    if 'idEvento' not in request.args and 'email' not in request.args:
        return "Error: No idEvento and email fields provided."
    elif 'idEvento' not in request.args:
        return "Error: No idEvento field provided. Please specify an id."
    elif 'email' not in request.args:
        return "Error: No email field provided. Please specify an email."

    idEvento = request.args['idEvento']
    email = request.args['email']
    dataInscricao = request.args['dataInscricao']

    query = "INSERT INTO participacoes ( idEvento, email, dataInscricao ) VALUES( '"+idEvento+"', '"+email+"', '"+dataInscricao+"' );"
    print(query)
    records = mysqlInsert(conn, query)
    mysqlCloseConnection(conn)
    print(records)
    return jsonify(records)

@app.route('/participacoes/remove', methods=['GET'])
def participacoesRemove():
    conn = mysqlConnection()
    if 'idEvento' not in request.args and 'email' not in request.args:
        return "Error: No idEvento and email fields provided."
    elif 'idEvento' not in request.args:
        return "Error: No idEvento field provided. Please specify an id."
    elif 'email' not in request.args:
        return "Error: No email field provided. Please specify an email."

    idEvento = request.args['idEvento']
    email = request.args['email']

    query = "DELETE FROM participacoes WHERE idEvento='"+idEvento+"' AND email='"+email+"';"
    records = mysqlUpdate(conn, query)
    mysqlCloseConnection(conn)
    return jsonify(records)

@app.route('/participacoes/update', methods=['GET'])
def participacoesUpdate():
    conn = mysqlConnection()
    if 'idEvento' not in request.args and 'email' not in request.args:
        return "Error: No idEvento and email fields provided."
    elif 'idEvento' not in request.args:
        return "Error: No idEvento field provided. Please specify an id."
    elif 'email' not in request.args:
        return "Error: No email field provided. Please specify an email."

    idEvento = request.args['idEvento']
    email = request.args['email']

    if 'dataInscricao' not in request.args:
        query = "SELECT dataInscricao FROM participacoes WHERE idEvento='"+idEvento+"' AND email='"+email+"';"
        records = mysqlQuery(conn, query)
        dataInscricao = str(records[0][0])
    else:
        dataInscricao = request.args['dataInscricao']

    query = "UPDATE participacoes SET dataInscricao='"+dataInscricao+"' WHERE idEvento='"+idEvento+"' AND email='"+email+"';"

    records = mysqlUpdate(conn, query)
    mysqlCloseConnection(conn)
    return jsonify(records)


import logging
from asyncio import (
    AbstractEventLoop,
)
from typing import (
    NoReturn,
)

from aiomisc.entrypoint import (
    Entrypoint,
)
from aiomisc.utils import (
    create_default_event_loop,
)
from cached_property import (
    cached_property,
)

logger = logging.getLogger(__name__)

class EntrypointLauncher:
    """EntryPoint Launcher class."""

    def __init__(self, services: tuple):
        self.services = services

    def launch(self) -> NoReturn:
        """Launch a new execution and keeps running forever..

        :return: This method does not return anything.
        """
        logger.info("Starting Discovery...")
        try:
            self.loop.run_until_complete(self.entrypoint.__aenter__())
            logger.info("Discovery is up and running!")
            self.loop.run_forever()
        except KeyboardInterrupt:  # pragma: no cover
            logger.info("Stopping discovery...")
        finally:
            self.graceful_shutdown()

    def graceful_shutdown(self, err: Exception = None) -> None:
        """Shutdown the services execution gracefully.

        :return: This method does not return anything.
        """
        self.loop.run_until_complete(self.entrypoint.graceful_shutdown(err))

    @cached_property
    def entrypoint(self) -> Entrypoint:
        """Entrypoint instance.

        :return: An ``Entrypoint`` instance.
        """

        return Entrypoint(*self.services, loop=self.loop)  # pragma: no cover

    @cached_property
    def loop(self) -> AbstractEventLoop:
        """Create the loop.

        :return: An ``AbstractEventLoop`` instance.
        """
        return create_default_event_loop()[0]  # pragma: no cover