Update app.py

app.py

@@ -83,14 +83,13 @@ def format_prompt(message, history):
     prompt += f"[INST] {message} [/INST]"
     return prompt
 
-def generate(prompt, history=[], temperature=0.1, max_new_tokens=25000, top_p=0.95, repetition_penalty=1.0):
+def generate(prompt, history=[], temperature=0.1, max_new_tokens=24000, top_p=0.95, repetition_penalty=1.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: 입력이 최대 허용 토큰 수를 초과합니다.", input_tokens
-        return
+        return "Error: 입력이 최대 허용 토큰 수를 초과합니다.", f"Used tokens: {input_tokens}"
 
     formatted_prompt = format_prompt(prompt, history)
     try:
@@ -102,9 +101,10 @@ def generate(prompt, history=[], temperature=0.1, max_new_tokens=25000, top_p=0.
                 output += response['generated_text']
             else:
                 output += str(response)
-        yield output, input_tokens + max_new_tokens
+        return output, f"Used tokens: {input_tokens + max_new_tokens}"
     except Exception as e:
-        yield f"Error: {str(e)}", 0
+        return f"Error: {str(e)}", "Used tokens: 0"
+
 
 
 
@@ -129,14 +129,13 @@ def update_output(result):
     output_text, used_tokens = result
     return output_text, f"Used tokens: {used_tokens}"
 
-demo = gr.Interface(
+demo = gr.ChatInterface(
     fn=generate,
-    inputs=[gr.Textbox(label="질문을 입력하세요", placeholder="여기에 질문을 입력하세요...", lines=2), gr.JSON(label="History", value=[])],
-    outputs=[gr.Markdown(), gr.Label()],
-    examples=examples,
+    chatbot=mychatbot,
     title="AIQ 코드파일럿: OpenLLM v1.12",
-    description="AIQ Codepilot과 상호작용해 보세요.",
-    post_process=update_output  # 사용된 토큰 수 업데이트를 위한 포스트 프로세스 함수
+    retry_btn=None,
+    undo_btn=None,
+    examples=examples    
 )
 
-demo.launch(show_api=False)
+demo.queue().launch(show_api=False)