RomanTeucher/python-coder

1import torch2from peft import PeftModel, PeftConfig3from transformers import AutoModelForCausalLM, AutoTokenizer4 5peft_model_id = f"RomanTeucher/PythonCoder"6config = PeftConfig.from_pretrained(peft_model_id)7model = AutoModelForCausalLM.from_pretrained(8    config.base_model_name_or_path,9    return_dict=True,10    load_in_8bit=True,11    device_map="auto",12)13tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)14 15# Load the Lora model16model = PeftModel.from_pretrained(model, peft_model_id)17 18 19def make_inference(instruction):20    batch = tokenizer(f"Below is an instruction, please create a python function based on that.\n\n### Instruction:\n{instruction} \n\n### Code:", return_tensors='pt')21 22    with torch.cuda.amp.autocast():23        output_tokens = model.generate(**batch, max_new_tokens=50)24 25    return tokenizer.decode(output_tokens[0], skip_special_tokens=True)26 27 28if __name__ == "__main__":29    # make a gradio interface30    import gradio as gr31 32    gr.Interface(33        make_inference,34        [35            gr.inputs.Textbox(lines=2, label="Instruction"),36        ],37        gr.outputs.Textbox(label="Code"),38        title="PythonCoder",39        description="PythonCoder is a generative model that generates python code for simple instructions.",40    ).launch()