Spaces:
Running
Running
| import torch | |
| import gradio as gr | |
| from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, logging | |
| checkpoint = "Salesforce/codegen-350M-mono" | |
| tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True) | |
| model = AutoModelForCausalLM.from_pretrained(checkpoint, cache_dir="models/", trust_remote_code=True, revision="main") | |
| def code_gen(text, max_tokens, temp, top_p, rep_penality): | |
| logging.set_verbosity(logging.CRITICAL) | |
| pipe = pipeline( | |
| model=checkpoint, | |
| max_new_tokens=max_tokens, | |
| temperature=temp, | |
| top_p=top_p, | |
| device= "cuda" if torch.cuda.is_available() else "cpu", | |
| repetition_penalty=rep_penality | |
| ) | |
| response = pipe(text) | |
| print(response) | |
| return response[0]['generated_text'] | |
| Inferece = gr.Interface( | |
| fn=code_gen, | |
| inputs=[ | |
| gr.components.Textbox(label="Enter your request, and the AI will generate the code for you."), | |
| gr.components.Slider(minimum=128, maximum=1024, step=128, value=512, label="Choose Max Token Size"), | |
| gr.components.Slider(minimum=0.1, maximum=1, step=0.05, value=0.65, label="Choose the model Temperature"), | |
| gr.components.Slider(minimum=0.1, maximum=1.25, step=0.05, value=0.9, label="Choose top_p"), | |
| gr.components.Slider(minimum=0.1, maximum=2, step=0.1, value=1.15, label="Choose repetition_penalty") | |
| ], | |
| outputs="text", | |
| title="AI Code Gen", | |
| live=False | |
| ) | |
| Inferece.queue(concurrency_count=1) | |
| Inferece.launch() |