Update README.md

README.md

@@ -21,7 +21,7 @@ pipeline_tag: text-generation
 </div>
 
 
-# Gemma Coder 🦙👩‍💻
+# Gemma Coder 👩‍💻
 **Gemma 2B** fine-tuned on the **CodeAlpaca 20k instructions dataset** by using the method **QLoRA** with [PEFT](https://github.com/huggingface/peft) library.
 
 ## Model description 🧠
@@ -95,15 +95,15 @@ def create_prompt(instruction):
 
 def generate(
         instruction,
-        max_new_tokens=128,
+        max_new_tokens=256,
         temperature=0.1,
         top_p=0.75,
         top_k=40,
-        num_beams=4,
+        num_beams=2,
         **kwargs,
 ):
-    prompt = create_prompt(instruction)
-    print(prompt)
+    system = f"<bos><|system|>\nYou are a helpful coding assistant.<eos>\n"
+    prompt = f"{system}<|user|>\n{instruction}<eos>\n<|assistant|>\n"
     inputs = tokenizer(prompt, return_tensors="pt")
     input_ids = inputs["input_ids"].to("cuda")
     attention_mask = inputs["attention_mask"].to("cuda")
@@ -120,13 +120,13 @@ def generate(
             attention_mask=attention_mask,
             generation_config=generation_config,
             return_dict_in_generate=True,
-            output_scores=True,
+            #output_scores=True,
             max_new_tokens=max_new_tokens,
             early_stopping=True
         )
     s = generation_output.sequences[0]
-    output = tokenizer.decode(s)
-    return output.split("### Solution:")[1].lstrip("\n")
+    output = tokenizer.decode(s, skip_special_tokens=True)
+    return output.split("<|assistant|>")[1]
 
 instruction = """
 Edit the following XML code to add a navigation bar to the top of a web page