Update app.py

app.py

@@ -6,14 +6,21 @@ from utils.prompter import Prompter
 
 class CustomPrompter(Prompter):
     def get_response(self, output: str) -> str:
-        return output.split(self.template["response_split"])[1].strip().split("### Instruction:")[0]
+        # Safely split on '### Response:'
+        split_output = output.split(self.template["response_split"], maxsplit=1)
+        if len(split_output) < 2:
+            return output.strip()
+        response_part = split_output[1].strip()
 
-prompt_template_name = "alpaca"  # The prompt template to use, will default to alpaca.
+        # Optionally strip out any subsequent '### Instruction:'
+        end_index = response_part.find("### Instruction:")
+        if end_index != -1:
+            response_part = response_part[:end_index].strip()
 
-prompter = CustomPrompter(prompt_template_name)
+        return response_part
 
-model = AutoModelForCausalLM.from_pretrained("Joaoffg/ELM")
-tokenizer = AutoTokenizer.from_pretrained("Joaoffg/ELM")
+prompt_template_name = "alpaca"
+prompter = CustomPrompter(prompt_template_name)
 
 def tokenize(prompt, add_eos_token=True):
   result = tokenizer(
@@ -62,52 +69,47 @@ def generate_and_tokenize_prompt(data_point):
     return tokenized_full_prompt
 
 def evaluate(instruction):
-    try:
-        # Generate a response:
-        input_text = None
-        prompt = prompter.generate_prompt(instruction, input_text)
-        inputs = tokenizer(prompt, return_tensors="pt")
-        input_ids = inputs["input_ids"]
-        
-        temperature = 0.2
-        top_p = 0.95
-        top_k = 25
-        num_beams = 1
-        max_new_tokens = 256
-        repetition_penalty = 1.3
-        do_sample = False
-        num_return_sequences = 1
-        
-        generation_config = transformers.GenerationConfig(
-            temperature=temperature,
-            top_p=top_p,
-            top_k=top_k,
-            num_beams=num_beams,
-            repetition_penalty=repetition_penalty,
-            do_sample=do_sample,
-            min_new_tokens=32,
-            num_return_sequences=num_return_sequences,
-            pad_token_id=0
+    input = None
+    prompt = prompter.generate_prompt(instruction, input)
+    inputs = tokenizer(prompt, return_tensors="pt")
+    input_ids = inputs["input_ids"]
+
+    # Example generation config
+    temperature=0.2
+    top_p=0.95
+    top_k=25
+    num_beams=1
+    max_new_tokens=256
+    repetition_penalty = 2.0
+    do_sample = True
+    
+    generation_config = transformers.GenerationConfig(
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k,
+        num_beams=num_beams,
+        repetition_penalty=repetition_penalty,
+        do_sample=do_sample,
+        min_new_tokens=32,
+        num_return_sequences=1,
+        pad_token_id=0,
+        # Optionally define a stopping criterion to stop at '### Instruction:'
+        # stopping_criteria=StoppingCriteriaList([StopOnTokens(tokenizer.encode("### Instruction:", add_special_tokens=False))]),
+    )
+
+    with torch.no_grad():
+        generation_output = model.generate(
+            input_ids=input_ids,
+            generation_config=generation_config,
+            return_dict_in_generate=True,
+            output_scores=True,
+            max_new_tokens=max_new_tokens,
         )
-        
-        with torch.no_grad():
-            generation_output = model.generate(
-                input_ids=input_ids,
-                generation_config=generation_config,
-                return_dict_in_generate=True,
-                output_scores=True,
-                max_new_tokens=max_new_tokens,
-            )
-        
-        print(f'Instruction: {instruction}')
-        
-        for i, s in enumerate(generation_output.sequences):
-            output = tokenizer.decode(s, skip_special_tokens=True)
-            return prompter.get_response(output)
-    except Exception as e:
-        return str(e)
-
-# Define the Gradio interface
+
+    # For demo, just take the first sequence
+    output = tokenizer.decode(generation_output.sequences[0], skip_special_tokens=True)
+    return prompter.get_response(output)
+
 interface = gr.Interface(
     fn=evaluate,
     inputs=[
@@ -132,5 +134,4 @@ interface = gr.Interface(
     ),
 )
 
-# Launch the Gradio interface
 interface.queue().launch()