Update app.py

app.py

@@ -1,7 +1,9 @@
 from huggingface_hub import InferenceClient
 import gradio as gr
+from transformers import GPT2Tokenizer
 
 client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
+tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
 
 # 시스템 인스트럭션을 설정하지만 사용자에게 노출하지 않습니다.
 system_instruction = """
@@ -74,39 +76,29 @@ const feedbackLink = `https://promptspellsmith.com/feedback`
 """
 
 def format_prompt(message, history):
-    prompt = "<s>"
-    # 시스템 인스트럭션을 프롬프트에 포함하지만, 이를 사용자에게는 표시하지 않습니다.
-    prompt += f"[SYSTEM] {system_instruction} [/SYSTEM]"
+    prompt = "<s>[SYSTEM] {} [/SYSTEM]".format(system_instruction)
     for user_prompt, bot_response in history:
-        prompt += f"[INST] {user_prompt} [/INST]"
-        prompt += f" {bot_response}</s> "
+        prompt += f"[INST] {user_prompt} [/INST]{bot_response}</s> "
     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):
-    temperature = float(temperature)
-    if temperature < 1e-2:
-        temperature = 1e-2
-    top_p = float(top_p)
-
-    generate_kwargs = dict(
-        temperature=temperature,
-        max_new_tokens=max_new_tokens,
-        top_p=top_p,
-        repetition_penalty=repetition_penalty,
-        do_sample=True,
-        seed=42,
-    )
-
+    # 토큰 수 계산
+    input_tokens = len(tokenizer.encode(prompt))
+    available_tokens = 32768 - input_tokens  # 최대 허용 토큰 수에서 입력 토큰 수를 뺀 값
+    max_new_tokens = min(max_new_tokens, available_tokens)  # 사용 가능한 최대 토큰 수 조정
+    
     formatted_prompt = format_prompt(prompt, history)
-
-    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
-    output = ""
-
-    for response in stream:
-        output += response.token.text
-        yield output
-    return output
+    
+    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)
+        output = ""
+        for response in stream:
+            output += response.token.text
+            yield output
+    except Exception as e:
+        yield f"Error: {str(e)}"
 
 mychatbot = gr.Chatbot(
     avatar_images=["./user.png", "./botm.png"],
@@ -116,23 +108,20 @@ mychatbot = gr.Chatbot(
     likeable=True,
 )
 
-
 examples = [
     ["좋은 예제를 알려줘."],
     ["requirements.txt 출력"],
     ["전체 코드를 다시 출력"],
     ["코드 오류를 확인하고 자세히 설명해줘."],
-    ["Huggingface와 Gradio를 사용하는 방법에 대해 물어보세요."]    
+    ["Huggingface와 Gradio를 사용하는 방법에 대해 물어보세요."]
 ]
 
-
-demo = gr.ChatInterface(
+demo = gr.Interface(
     fn=generate,
-    chatbot=mychatbot,
-    title="AIQ 코드파일럿: OpenLLM v1.12",
-    retry_btn=None,
-    undo_btn=None,
-    examples=examples    
+    inputs=gr.Textbox(placeholder="여기에 질문을 입력하세요...", lines=2),
+    outputs=gr.Markdown(),
+    examples=examples,
+    title="AIQ 코드파일럿: OpenLLM v1.12"
 )
 
-demo.queue().launch(show_api=False)
+demo.launch(show_api=False)