PipableAI
/

pip-code-bandit

Model card Files Files and versions Community

pratham-pi commited on May 13, 2024

Commit

ca8b07f

·

verified ·

1 Parent(s): aceac3b

Update README.md

Files changed (1) hide show

README.md +5 -4

README.md CHANGED Viewed

@@ -153,13 +153,14 @@ prompt = """
 <question>
 Generate a python function for adding two numbers.
 </question>
-<response>
 """
 inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
 outputs = model.generate(**inputs, max_new_tokens=new_tokens)
-out = (
-    tokenizer.decode(outputs[0][inputs.input_ids.shape[-1]:], skip_special_tokens=True)
-)
 ```

 <question>
 Generate a python function for adding two numbers.
 </question>
+<code>
 """
 inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
 outputs = model.generate(**inputs, max_new_tokens=new_tokens)
+response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+response = response.split("<code>")[1].split("</code>")[0]
+print(response)
 ```