Update app.py

app.py

@@ -28,14 +28,14 @@ with open("model_configs.json", "r") as f:
 # Extract instruction
 extract_input = model_config["extract_input"]
 
-terminators = [
-    tokenizer.eos_token_id,
-    tokenizer.convert_tokens_to_ids("<|eot_id|>"),
-]
-
 
 @spaces.GPU
-def generate_instruction():
+def generate_instruction_response():
+    terminators = [
+        tokenizer.eos_token_id,
+        tokenizer.convert_tokens_to_ids("<|eot_id|>"),
+    ]
+
     instruction = pipeline(
         extract_input,
         max_new_tokens=2048,
@@ -45,11 +45,13 @@ def generate_instruction():
         top_p=1,
     )
 
-    return instruction[0]["generated_text"][len(extract_input) :].split("\n")[0]
+    sanitized_instruction = instruction[0]["generated_text"][
+        len(extract_input) :
+    ].split("\n")[0]
 
+    response_template = f"""<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n{sanitized_instruction}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"""
 
-def generate_response(response_template):
-    return pipeline(
+    response = pipeline(
         response_template,
         max_new_tokens=2048,
         eos_token_id=terminators,
@@ -58,13 +60,7 @@ def generate_response(response_template):
         top_p=1,
     )
 
-
-def generate_instruction_response():
-    sanitized_instruction = generate_instruction()
-    response_template = f"""<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n{sanitized_instruction}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"""
-
     user_message = sanitized_instruction
-    response = generate_response(response_template)
     assistant_response = response[0]["generated_text"][len(response_template) :]
 
     return user_message, assistant_response