Update generator.py

generator.py

@@ -1,7 +1,18 @@
 from transformers import pipeline
 
 
+
+def format(self, context: str) -> str:
+    system_message = f"You are an expert in writing novels. You have won many prices."
+    user_message = f"Please assist based on the following context: {context}"
+    prompt = f"<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n{system_message}<|eot_id|>" 
+    prompt += f"<|start_header_id|>user<|end_header_id|>\n\n{user_input}<|eot_id|>"
+    prompt += "<|start_header_id|>assistant<|end_header_id|>\n\n"
+    return prompt
+        
 def generate(text, the_model, max_length, temperature, repetition_penalty):
+    print(text)
+    text = format(text)
     generator = pipeline("text-generation", model=the_model)
     result = generator(
         text,