Update main.py

main.py

@@ -36,32 +36,42 @@ class InputData(BaseModel):
     max_new_tokens: int = 2000
     top_p: float = 0.95
     repetition_penalty: float = 1.0
-
-class InputDataAsync(InputData):
+    Async: bool = false
     NumeroGenerazioni: int = 1
     StringaSplit: str = '********'
     NumeroCaratteriSplitInstruction: int = 30000
+
+class InputDataAsync(InputData):
     
 class PostSpazio(BaseModel):
     nomeSpazio: str
     input: str = '' 
     api_name: str = "/chat"
+
+def LoggaTesto(type, text)
+    print(f"{datetime.now()}: ----------------------------------| {type} |-----------------------------------\n{text}\n\n")
     
 #--------------------------------------------------- Generazione TESTO ------------------------------------------------------
 @app.post("/Genera")
 def generate_text(request: Request, input_data: InputData):
-    temperature = input_data.temperature
-    max_new_tokens = input_data.max_new_tokens
-    top_p = input_data.top_p
-    repetition_penalty = input_data.repetition_penalty
-    input_text = generate_input_text(input_data)  
-    max_new_tokens = min(max_new_tokens, 29500 - len(input_text))
-    print(f"{datetime.now()} - Input Text: {input_text}")
-    history = []
-    generated_response = generate(input_text, history, temperature, max_new_tokens, top_p, repetition_penalty)
-    print(f"{datetime.now()} - Response Text: {generated_response}")
-    return {"response": generated_response}
-
+    LoggaTesto("INPUT", input_data.input)
+    if not input_data.Async: 
+        temperature = input_data.temperature
+        max_new_tokens = input_data.max_new_tokens
+        top_p = input_data.top_p
+        repetition_penalty = input_data.repetition_penalty
+        input_text = generate_input_text(input_data)  
+        max_new_tokens = min(max_new_tokens, 29500 - len(input_text))
+        history = []
+        generated_response = generate(input_text, history, temperature, max_new_tokens, top_p, repetition_penalty)
+        LoggaTesto("RISPOSTA SINCRONA", generated_response)
+        return {"response": generated_response}
+    else: 
+        input_data.Async = false
+        result_data = asyncio.run(GeneraTestoAsync("https://matteoscript-fastapi.hf.space/Genera", input_data))
+        LoggaTesto("RISPOSTA ASINCRONA FINALE", result_data)
+        return {"response": result_data}
+    
 def generate_input_text(input_data):
     if input_data.instruction.startswith("http"):
         try:
@@ -124,7 +134,6 @@ def format_prompt(message, history):
 #--------------------------------------------------- Generazione TESTO ASYNC ------------------------------------------------------
 @app.post("/GeneraAsync")
 def generate_textAsync(request: Request, input_data: InputDataAsync):
-    print(input_data.input)
     result_data = asyncio.run(GeneraTestoAsync("https://matteoscript-fastapi.hf.space/Genera", input_data))
     return {"response": result_data}
 
@@ -144,13 +153,10 @@ async def make_request(session, token, data, url, max_retries=3):
                 print(result_data)
                 return result_data
         except (asyncio.TimeoutError, aiohttp.ClientError, requests.exceptions.HTTPError) as e:
-            print(f"Error: {e}.")
-            
+            LoggaTesto("ERRORE ASYNC", {e})
             if isinstance(e, (asyncio.TimeoutError, requests.exceptions.HTTPError)) and e.response.status in [502, 504]:
-                print("Skipping retries for this error.")
                 break
 
-            print("Retrying...")
             await asyncio.sleep(1)  
     raise Exception("Max retries reached or skipping retries. Unable to make the request.")
 
@@ -165,7 +171,6 @@ async def make_request_old(session, token, data, url):
             result_data = await response.json()
         except aiohttp.ContentTypeError:
             result_data = await response.text()   
-        print(result_data)
         return result_data
 
 async def CreaListaInput(input_data):
@@ -176,7 +181,6 @@ async def CreaListaInput(input_data):
             input_data.instruction = resp.text
         except requests.exceptions.RequestException as e:
             input_data.instruction = ""
-    print(input_data.instruction)
     try:
         lista_dizionari = []
         nuova_lista_dizionari = []
@@ -187,7 +191,6 @@ async def CreaListaInput(input_data):
             nuova_lista_dizionari = DividiInstructionText(input_data)
     except json.JSONDecodeError:
         nuova_lista_dizionari = DividiInstructionText(input_data)
-    print(nuova_lista_dizionari)    
     return nuova_lista_dizionari
 
 def split_at_space_or_dot(input_string, length):
@@ -196,7 +199,6 @@ def split_at_space_or_dot(input_string, length):
     valid_positions = [pos for pos in positions if pos >= 0]
     lastpos = max(valid_positions) if valid_positions else length
     indice_divisione = int(lastpos)
-    print(indice_divisione)
     return indice_divisione + 1
 
 def DividiInstructionJSON(lista_dizionari, input_data):
@@ -271,8 +273,7 @@ async def GeneraTestoAsync(url, input_data):
         tasks = []
         ListaInput = await CreaListaInput(input_data)
         for data in ListaInput:
-            print('------------------------------------------------------------------------------------------------')
-            print(data)
+            LoggaTesto("RICHIESTA ASINCRONA", data)
             tasks.extend([make_request(session, token, data, url) for _ in range(input_data.NumeroGenerazioni)])
         return await asyncio.gather(*tasks)