app.py

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=500
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 = "Gradio demo for Genji-python-6b: Genji is a transformer model finetuned on EleutherAI's GPT-J 6B model. This particular model is trained on python only code approaching 4GB in size. 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", lines=5), 
    gr.outputs.Textbox(label="Output text"),
    title=title,
    description=description,
    article=article,
    examples=[
        ['def print_customer_name']
    ]).launch(debug=True)