Update app.py

app.py

@@ -16,32 +16,32 @@ models = {
 def get_client(model):
     return InferenceClient(model=model, token=hf_token)
 
-def respond(message, chat_history, system_message, max_tokens, temperature, top_p, selected_model):
+def respond(message, history, system_message, max_tokens, temperature, top_p, selected_model):
     stop_event.clear()
-    client = get_client(selected_model)
+    client = InferenceClient(model=selected_model, token=hf_token)
     
-    messages = [{"role": "system", "content": system_message}]
-    messages.extend([{"role": "user" if i % 2 == 0 else "assistant", "content": m} for h in chat_history for i, m in enumerate(h) if m])
+    messages = [{"role": "system", "content": system_message + "\n사용자의 입력에만 직접적으로 답변하세요. 추가 질문을 생성하거나 사용자의 입력을 확장하지 마세요."}]
+    messages.extend([{"role": "user" if i % 2 == 0 else "assistant", "content": m} for h in history for i, m in enumerate(h) if m])
     messages.append({"role": "user", "content": message})
     
-    response = ""
     try:
-        for chunk in client.chat_completion(
-            messages,
-            max_tokens=max_tokens,
+        response = ""
+        for chunk in client.text_generation(
+            prompt="\n".join([f"{m['role']}: {m['content']}" for m in messages]),
+            max_new_tokens=max_tokens,
             temperature=temperature,
             top_p=top_p,
             stream=True
         ):
             if stop_event.is_set():
                 break
-            if hasattr(chunk.choices[0], 'delta'):
-                token = chunk.choices[0].delta.content
-            else:
-                token = chunk.choices[0].text
-            if token:
-                response += token
+            if chunk:
+                response += chunk
+                # 응답의 시작 부분에서 사용자 입력의 반복을 제거
+                if response.startswith(message):
+                    response = response[len(message):].lstrip()
                 yield response
+        
     except Exception as e:
         yield f"오류 발생: {str(e)}"