from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    GPTNeoForCausalLM,
)
import torch
import psutil
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-2.7B")

model = AutoModelForCausalLM.from_pretrained("NovelAI/genji-python-6B").half().eval().cuda()

import gradio as gr

maxLength=200
temperature=0.4
top_k = 50
top_p = 0.9
repetition_penalty = 1.13
repetition_penalty_range = 512
repetition_penalty_slope = 3.33
def generator(text):
    tokens = tokenizer(text, return_tensors="pt").input_ids.cuda()[:, -(2047-maxLength):]
    out = model.generate(
        tokens.long(),
        do_sample=True,
        min_length=tokens.shape[1] + maxLength,
        max_length=tokens.shape[1] + maxLength,
        temperature=temperature,
        top_k = top_k,
        top_p = top_p,
        repetition_penalty = repetition_penalty,
        repetition_penalty_range = repetition_penalty_range,
        repetition_penalty_slope = repetition_penalty_slope,
        use_cache=True,
        bad_words_ids=None,
        pad_token_id=tokenizer.eos_token_id,
    ).long().to("cpu")[0]
    return tokenizer.decode(out[tokens.shape[1]:])

title = "genji-python-6b"
description = "demo for Genji-python-6b. To use it, simply add your text, or click one of the examples to load them. Read more at the links below."
article = "<p style='text-align: center'><a href='https://colab.research.google.com/drive/1PnWpx02IEUkY8jhLKd_NewUGEXahAska'>Colab</a> | <a href='https://huggingface.co/NovelAI/genji-python-6B'>Huggingface Model</a></p>"

gr.Interface(
    generator, 
    [gr.inputs.Textbox(label="input text")], 
    gr.outputs.Textbox(label="Output text"),
    title=title,
    description=description,
    article=article,
    examples=[
        ['def print_customer_name']
    ]).launch(debug=True)