Update main.py

main.py

@@ -1,5 +1,5 @@
 from fastapi import FastAPI, HTTPException
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional
 from pydantic import BaseModel
 from os import getenv
 from huggingface_hub import InferenceClient
@@ -18,11 +18,11 @@ HF_TOKEN = getenv("HF_TOKEN")
 
 class InputData(BaseModel):
     model: str
-    system_prompt_template: str
-    prompt_template: str
+    system_prompt_template: List[str]
+    prompt_template: List[str]
     end_token: str
-    system_prompt: str
-    user_input: str
+    system_prompt: List[str]
+    user_input: List[str]
     history: str = ""
     segment: bool = False
     max_sentences: Optional[int] = None
@@ -36,24 +36,29 @@ 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:
-        data.history += data.prompt_template.replace("{Prompt}", data.user_input)
+        for prompt in data.prompt_template:
+            for user_input in data.user_input:
+                data.history += prompt.replace("{Prompt}", user_input) + "\n"
         return {
             "response": "",
             "history": data.history + data.end_token
         }
 
+    user_input_str = ""
     if data.segment:
-        user_sentences = tokenizer.tokenize(data.user_input)
-        user_input_str = "\n".join(user_sentences)
+        for user_input in data.user_input:
+            user_sentences = tokenizer.tokenize(user_input)
+            user_input_str += "\n".join(user_sentences) + "\n"
     else:
-        user_input_str = data.user_input
+        user_input_str = "\n".join(data.user_input)
 
-    data.history += data.prompt_template.replace("{Prompt}", user_input_str)
+    for prompt in data.prompt_template:
+        data.history += prompt.replace("{Prompt}", user_input_str) + "\n"
 
-    inputs = (
-        data.system_prompt_template.replace("{SystemPrompt}", data.system_prompt) +
-        data.history
-    )
+    inputs = ""
+    for system_prompt in data.system_prompt_template:
+        inputs += system_prompt.replace("{SystemPrompt}", "\n".join(data.system_prompt)) + "\n"
+    inputs += data.history
 
     seed = random.randint(0, 2**32 - 1)
 
@@ -111,4 +116,4 @@ async def check_word(data: WordCheckData) -> Dict[str, Any]:
         "found": found
     }
 
-    return result
+    return result