Update app.py

app.py

@@ -24,7 +24,7 @@ try:
     import psutil
     import resource
     total_memory = psutil.virtual_memory().total
-    limit = int(total_memory * 80.0)  # 1% del total en bytes # Corrección: Usar 0.01 para 1%
+    limit = int(total_memory * 0.01)  # 1% del total en bytes # Corrección: Usar 0.01 para 1%
     resource.setrlimit(resource.RLIMIT_AS, (limit, limit))
     print(f"Memory limit set to {limit} bytes (1% of total system memory).") # Imprimir para verificar el límite aplicado
 except Exception as e:
@@ -132,6 +132,7 @@ async def generate(request: GenerateRequest):
             repetition_penalty=repetition_penalty,
             do_sample=do_sample,
             num_return_sequences=num_return_sequences,
+            stream=stream, # Add stream=True/False to generation config
         )
 
         stop_token_ids = []
@@ -160,7 +161,6 @@ async def stream_text(model, tokenizer, input_text, generation_config, stopping_
     """
     # Limitar la entrada para minimizar el uso de memoria
     encoded_input = tokenizer(input_text, return_tensors="pt", truncation=True, max_length=max_length).to(device)
-    encoded_input_len = encoded_input["input_ids"].shape[-1]
 
     # Con torch.no_grad() se evita almacenar información para gradientes
     with torch.no_grad():
@@ -169,18 +169,15 @@ async def stream_text(model, tokenizer, input_text, generation_config, stopping_
             **encoded_input,
             generation_config=generation_config,
             stopping_criteria=stopping_criteria_list,
-            return_dict_in_generate=True,
-            output_scores=True,
-            # stream=True, # Eliminar 'stream=True' aquí, ya que GenerationConfig lo maneja
+            # return_dict_in_generate=True, # Remove return_dict_in_generate for streaming
+            # output_scores=True, # output_scores might not be needed for streaming text only
         ):
-            # Se extraen solo los tokens generados (excluyendo la entrada)
-            new_tokens = output.sequences[:, encoded_input_len:]
-            for token_batch in new_tokens:
-                token = tokenizer.decode(token_batch, skip_special_tokens=True)
-                if token:
-                    # Se envía cada token inmediatamente
-                    yield token
-                    await asyncio.sleep(chunk_delay)
+            # In streaming mode, output is directly the generated token IDs
+            token = tokenizer.decode(output, skip_special_tokens=True)
+            if token:
+                # Se envía cada token inmediatamente
+                yield token
+                await asyncio.sleep(chunk_delay)
     await cleanup_memory(device)
 
 async def generate_non_stream(model, tokenizer, input_text, generation_config, stopping_criteria_list, device, max_length=64):