modified: app.py

app.py

@@ -13,8 +13,12 @@ model = AutoModelForCausalLM.from_pretrained(checkpoint, trust_remote_code=True,
 def generate(text, max_length, num_return_sequences=1):
     inputs = tokenizer.encode(text, padding=False, add_special_tokens=False, return_tensors="pt")
     outputs = model.generate(inputs, max_length=max_length, num_return_sequences=num_return_sequences)
-    gen_text = "Assignment : " + tokenizer.decode(outputs[0]).split("#")[0] if "#" else "Assignment : " + tokenizer.decode(outputs[0])
-    return gen_text
+    gen_text = "Assignment : " + tokenizer.decode(outputs[0])
+    if gen_text.count("#") > 2:
+      split_text = gen_text.split("#", 2)
+      return split_text[0] + "#" + split_text[1]
+    else:
+      return gen_text
 
 
 def extract_functions(text):
@@ -58,6 +62,7 @@ with demo:
     examples = [
     ["generate a python for sum number"],
     ["generate a python function to find max min element of list"],
+    ["generate a python function to find minimum of two numbers with test case"],
     ]
 
     gr.Examples(examples=examples, inputs=inputs, cache_examples=False)