lvwerra HF staff commited on
Commit
fd41877
1 Parent(s): 95ef04e

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +53 -0
app.py ADDED
@@ -0,0 +1,53 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed
3
+ from transformers import pipeline
4
+ import os
5
+
6
+ description = """# SantaCoder Endpoint"""
7
+ token = os.environ["HUB_TOKEN"]
8
+ device="cuda:0"
9
+
10
+ tokenizer = AutoTokenizer.from_pretrained("bigcode/christmas-models", use_auth_token=token)
11
+ model = AutoModelForCausalLM.from_pretrained("bigcode/christmas-models", trust_remote_code=True, use_auth_token=token)
12
+
13
+
14
+ def code_generation(gen_prompt, max_tokens, temperature=0.6, seed=42):
15
+ set_seed(seed)
16
+ pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
17
+ generated_text = pipe(gen_prompt, do_sample=True, top_p=0.95, temperature=temperature, max_new_tokens=max_tokens)[0]['generated_text']
18
+ return generated_text
19
+
20
+
21
+ demo = gr.Blocks()
22
+ with demo:
23
+ with gr.Row():
24
+ gr.Markdown(value=description)
25
+ with gr.Row():
26
+ with gr.Column():
27
+ code = gr.Textbox(lines=10, label="Input code")
28
+ max_tokens= gr.Slider(
29
+ minimum=8,
30
+ maximum=1000,
31
+ step=1,
32
+ label="Number of tokens to generate",
33
+ )
34
+ temperature = gr.Slider(
35
+ minimum=0.1,
36
+ maximum=2.5,
37
+ step=0.1,
38
+ label="Temperature",
39
+ )
40
+ seed = gr.Slider(
41
+ minimum=0,
42
+ maximum=1000,
43
+ step=1,
44
+ label="Random seed to use for the generation"
45
+ )
46
+ run = gr.Button()
47
+ with gr.Column():
48
+ output = gr.Textbox(lines=10, label="Generated code")
49
+
50
+ event = run.click(code_generation, [code, max_tokens, temperature, seed], output)
51
+ gr.HTML(label="Contact", value="<img src='https://huggingface.co/datasets/bigcode/admin/resolve/main/bigcode_contact.png' alt='contact' style='display: block; margin: auto; max-width: 800px;'>")
52
+
53
+ demo.launch()