Update main.py

main.py

@@ -1,5 +1,5 @@
 from fastapi import FastAPI, HTTPException
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 from pydantic import BaseModel
 from os import getenv
 from huggingface_hub import InferenceClient
@@ -24,6 +24,8 @@ class InputData(BaseModel):
     system_prompt: str
     user_input: str
     history: str = ""
+    segment: bool = False
+    max_sentences: Optional[int] = None
 
 class WordCheckData(BaseModel):
     string: str
@@ -31,12 +33,21 @@ class WordCheckData(BaseModel):
 
 @app.post("/generate-response/")
 async def generate_response(data: InputData) -> Dict[str, Any]:
-    client = InferenceClient(model=data.model, token=HF_TOKEN)
+    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)
+        return {
+            "response": "",
+            "history": data.history + data.end_token
+        }
 
-    user_sentences = tokenizer.tokenize(data.user_input)
+    if data.segment:
+        user_sentences = tokenizer.tokenize(data.user_input)
+        user_input_str = "\n".join(user_sentences)
+    else:
+        user_input_str = data.user_input
 
-    user_input_str = "\n".join(user_sentences)
-    
     data.history += data.prompt_template.replace("{Prompt}", user_input_str)
 
     inputs = (
@@ -47,6 +58,7 @@ async def generate_response(data: InputData) -> Dict[str, Any]:
     seed = random.randint(0, 2**32 - 1)
 
     try:
+        client = InferenceClient(model=data.model, token=HF_TOKEN)
         response = client.text_generation(
             inputs,
             temperature=1.0,
@@ -56,17 +68,21 @@ async def generate_response(data: InputData) -> Dict[str, Any]:
 
         response_str = str(response)
 
-        ai_sentences = tokenizer.tokenize(response_str)
+        if data.segment:
+            ai_sentences = tokenizer.tokenize(response_str)
+            if data.max_sentences is not None:
+                ai_sentences = ai_sentences[:data.max_sentences]
+            ai_response_str = "\n".join(ai_sentences)
+        else:
+            ai_response_str = response_str
+
+        data.history += ai_response_str + "\n"
 
         cleaned_response = {
-            "New response": ai_sentences,
-            "Sentence count": min(len(ai_sentences), 3)
+            "New response": ai_sentences if data.segment else [response_str],
+            "Sentence count": len(ai_sentences) if data.segment else 1
         }
 
-        ai_response_str = "\n".join(ai_sentences)
-
-        data.history += ai_response_str + "\n"
-
         return {
             "response": cleaned_response,
             "history": data.history + data.end_token