Update app.py

app.py

@@ -87,13 +87,12 @@ def generate(prompt, history=[], temperature=0.1, max_new_tokens=25000, top_p=0.
     input_tokens = len(tokenizer.encode(prompt))
     available_tokens = 32768 - input_tokens
     max_new_tokens = min(max_new_tokens, available_tokens)
-    
+
     if available_tokens <= 0:
-        yield "Error: 입력이 최대 허용 토큰 수를 초과합니다."
+        yield "Error: 입력이 최대 허용 토큰 수를 초과합니다.", input_tokens
         return
 
     formatted_prompt = format_prompt(prompt, history)
-    
     try:
         stream = client.text_generation(formatted_prompt, temperature=temperature, max_new_tokens=max_new_tokens,
                                         top_p=top_p, repetition_penalty=repetition_penalty, do_sample=True, seed=42, stream=True)
@@ -102,10 +101,11 @@ def generate(prompt, history=[], temperature=0.1, max_new_tokens=25000, top_p=0.
             if isinstance(response, dict) and 'generated_text' in response:
                 output += response['generated_text']
             else:
-                output += str(response)  # 예외 처리를 개선하여 모든 응답을 문자열로 처리
-            yield output
+                output += str(response)
+        yield output, input_tokens + max_new_tokens
     except Exception as e:
-        yield f"Error: {str(e)}  # 에러 로깅을 강화하여 문제를 좀 더 명확히 식별"
+        yield f"Error: {str(e)}", 0
+
 
 
 mychatbot = gr.Chatbot(
@@ -125,14 +125,18 @@ examples = [
     ["Huggingface와 Gradio를 사용하는 방법에 대해 물어보세요.", []]
 ]
 
+def update_output(result):
+    output_text, used_tokens = result
+    return output_text, f"Used tokens: {used_tokens}"
 
-demo = gr.ChatInterface(
+demo = gr.Interface(
     fn=generate,
-    chatbot=mychatbot,
+    inputs=[gr.Textbox(label="질문을 입력하세요", placeholder="여기에 질문을 입력하세요...", lines=2), gr.JSON(label="History", value=[])],
+    outputs=[gr.Markdown(), gr.Label()],
+    examples=examples,
     title="AIQ 코드파일럿: OpenLLM v1.12",
-    retry_btn=None,
-    undo_btn=None,
-    examples=examples    
+    description="AIQ Codepilot과 상호작용해 보세요.",
+    post_process=update_output  # 사용된 토큰 수 업데이트를 위한 포스트 프로세스 함수
 )
 
-demo.queue().launch(show_api=False)
+demo.launch(show_api=False)