update app

app.py

@@ -6,31 +6,34 @@ from transformers import pipeline
 title = "InCoder Generator"
 description = "This is a subspace to make code generation with [InCoder](https://huggingface.co/facebook/incoder-6B), it is used in a larger [space](https://huggingface.co/spaces/loubnabnl/Code-generation-models-v1) for model comparison."
 example = [
-    ["def print_hello_world():", "Sample", 8, 42],
-    ["def get_file_size(filepath):", "Sample", 22, 42]]
+    ["def count_words(filename):", 40, 0.6, 42],
+    ["def print_hello_world():", "Sample", 8, 0.6, 42],
+    ["def get_file_size(filepath):", "Sample", 22, 0.6, 42]]
 tokenizer = AutoTokenizer.from_pretrained("facebook/incoder-1B")
 model = AutoModelForCausalLM.from_pretrained("facebook/incoder-1B", low_cpu_mem_usage=True)
     
 
-def code_generation(gen_prompt, strategy, max_tokens, seed=42):
+MAX_LENGTH = 2048
+BOS = "<|endoftext|>"
+def generate(gen_prompt, max_tokens, temperature=0.6, seed=42):
     set_seed(seed)
-    gen_kwargs = {}
-    gen_kwargs["do_sample"] = strategy == "Sample"
-    gen_kwargs["max_new_tokens"] = max_tokens
-    if gen_kwargs["do_sample"]:
-        gen_kwargs["temperature"] = 0.2
-        gen_kwargs["top_k"] = 0
-        gen_kwargs["top_p"] = 0.95
-    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
-    generated_text = pipe(gen_prompt, **gen_kwargs)[0]['generated_text']
+    input_ids = tokenizer(gen_prompt, return_tensors="pt").input_ids
+    current_length = input_ids.flatten().size(0)
+    max_length = max_tokens + current_length
+    if max_length > MAX_LENGTH:
+        max_length = MAX_LENGTH
+    output = model.generate(input_ids=input_ids, do_sample=True, top_p=0.95, temperature=temperature, max_length=max_length)
+    generated_text = tokenizer.decode(output.flatten())
+    if generated_text.startswith(BOS):
+        generated_text = detok_hypo_str[len(BOS):]
     return generated_text
 
+generation = generate(gen_prompt, length_limit=40, temperature=0.6)
 
 iface = gr.Interface(
-    fn=code_generation, 
+    fn=generate, 
     inputs=[
         gr.Textbox(lines=10, label="Input code"),
-        gr.Dropdown(choices=["Greedy", "Sample"], value="Sample"),
         gr.inputs.Slider(
             minimum=8,
             maximum=256,
@@ -38,6 +41,13 @@ iface = gr.Interface(
             default=8,
             label="Number of tokens to generate",
         ),
+        gr.inputs.Slider(
+            minimum=0,
+            maximum=2,
+            step=0.1,
+            default=0.6,
+            label="Temperature",
+        )
         gr.inputs.Slider(
             minimum=0,
             maximum=1000,