Update app.py

app.py

@@ -706,13 +706,12 @@ def stream_chat(
 
         print(f"입력 텐서 생성 후 CUDA 메모리: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
 
-        
-    try:
-        # 스트리머 초기화 시 타임아웃을 더 길게 설정
+        # 스트리머 설정
         streamer = TextIteratorStreamer(
             tokenizer, timeout=30.0, skip_prompt=True, skip_special_tokens=True
         )
-    
+        
+        # 생성 매개변수 설정
         generate_kwargs = dict(
             **inputs,
             streamer=streamer,
@@ -727,8 +726,6 @@ def stream_chat(
             use_cache=True
         )
 
-
-
         # 메모리 정리
         clear_cuda_memory()
 
@@ -736,26 +733,39 @@ def stream_chat(
         thread = Thread(target=model.generate, kwargs=generate_kwargs)
         thread.start()
 
-        # 스트리밍 처리 중 예외 처리 강화
+        # 응답 스트리밍
         buffer = ""
+        partial_message = ""
+        last_yield_time = time.time()
+        
         try:
             for new_text in streamer:
                 try:
                     buffer += new_text
-                    yield "", history + [[message, buffer]]
+                    partial_message += new_text
+                    
+                    # 일정 시간마다 또는 텍스트가 쌓일 때마다 결과 업데이트
+                    current_time = time.time()
+                    if current_time - last_yield_time > 0.1 or len(partial_message) > 20:
+                        yield "", history + [[message, buffer]]
+                        partial_message = ""
+                        last_yield_time = current_time
                 except Exception as inner_e:
                     print(f"개별 토큰 처리 중 오류: {str(inner_e)}")
                     continue
-        except Exception as stream_e:
-            print(f"스트리밍 전체 오류: {str(stream_e)}")
-            if not buffer:
-                buffer = "응답 생성 중 오류가 발생했습니다."
+                    
+            # 마지막 응답 확인
+            if buffer:
+                yield "", history + [[message, buffer]]
+                
+            # 대화 기록에 저장
+            chat_history.add_conversation(message, buffer)
+                
+        except Exception as e:
+            print(f"스트리밍 중 오류 발생: {str(e)}")
+            if not buffer:  # 버퍼가 비어있으면 오류 메시지 표시
+                buffer = f"응답 생성 중 오류가 발생했습니다: {str(e)}"
             yield "", history + [[message, buffer]]
-    except Exception as outer_e:
-        print(f"전체 생성 과정 오류: {str(outer_e)}")
-        yield "", history + [[message, "죄송합니다, 응답을 생성할 수 없습니다."]]
-    
-]
             
         # 스레드가 여전히 실행 중이면 종료 대기
         if thread.is_alive():