Update main.py

main.py

@@ -18,11 +18,11 @@ HF_TOKEN = getenv("HF_TOKEN")
 
 class InputData(BaseModel):
     model: str
-    system_prompt_template: List[str]
-    prompt_template: List[str]
+    system_prompt_template: str
+    prompt_template: str
     end_token: str
-    system_prompt: List[str]
-    user_input: List[str]
+    system_prompts: List[str]
+    user_inputs: List[str]
     history: str = ""
     segment: bool = False
     max_sentences: Optional[int] = None
@@ -36,28 +36,25 @@ async def generate_response(data: InputData) -> Dict[str, Any]:
     if data.max_sentences is not None and data.max_sentences != 0:
         data.segment = True
     elif data.max_sentences == 0:
-        for prompt in data.prompt_template:
-            for user_input in data.user_input:
-                data.history += prompt.replace("{Prompt}", user_input) + "\n"
+        for user_input in data.user_inputs:
+            data.history += data.prompt_template.replace("{Prompt}", user_input)
         return {
             "response": "",
             "history": data.history + data.end_token
         }
 
-    user_input_str = ""
     if data.segment:
-        for user_input in data.user_input:
+        for user_input in data.user_inputs:
             user_sentences = tokenizer.tokenize(user_input)
-            user_input_str += "\n".join(user_sentences) + "\n"
+            user_input_str = "\n".join(user_sentences)
+            data.history += data.prompt_template.replace("{Prompt}", user_input_str) + "\n"
     else:
-        user_input_str = "\n".join(data.user_input)
-
-    for prompt in data.prompt_template:
-        data.history += prompt.replace("{Prompt}", user_input_str) + "\n"
+        for user_input in data.user_inputs:
+            data.history += data.prompt_template.replace("{Prompt}", user_input) + "\n"
 
     inputs = ""
-    for system_prompt in data.system_prompt_template:
-        inputs += system_prompt.replace("{SystemPrompt}", "\n".join(data.system_prompt)) + "\n"
+    for system_prompt in data.system_prompts:
+        inputs += data.system_prompt_template.replace("{SystemPrompt}", system_prompt) + "\n"
     inputs += data.history
 
     seed = random.randint(0, 2**32 - 1)
@@ -116,4 +113,4 @@ async def check_word(data: WordCheckData) -> Dict[str, Any]:
         "found": found
     }
 
-    return result
+    return result