Update app.py

app.py

@@ -5,6 +5,43 @@ import os
 
 hf_token = os.getenv("HF_TOKEN")
 
+def get_model_response(client, messages, max_tokens, temperature, top_p, model_name):
+    prompt = "\n".join([f"{m['role']}: {m['content']}" for m in messages])
+    
+    try:
+        # 먼저 chat_completion 시도
+        response = client.chat_completion(
+            messages,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            stream=True
+        )
+        for message in response:
+            if hasattr(message.choices[0], 'delta'):
+                token = message.choices[0].delta.content
+            else:
+                token = message.choices[0].text
+            if token:
+                yield token
+    except Exception as chat_error:
+        try:
+            # chat_completion 실패 시 text_generation 시도
+            response = client.text_generation(
+                prompt,
+                max_new_tokens=max_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                stream=True
+            )
+            for token in response:
+                yield token
+        except Exception as text_error:
+            # 두 방법 모두 실패 시 오류 메시지 반환
+            yield f"모델 {model_name}에 대한 추론 실패:\n"
+            yield f"Chat 오류: {str(chat_error)}\n"
+            yield f"Text 오류: {str(text_error)}"
+
 def respond(message, history, system_message, max_tokens, temperature, top_p, selected_model):
     try:
         client = InferenceClient(model=selected_model, token=hf_token)
@@ -18,28 +55,17 @@ def respond(message, history, system_message, max_tokens, temperature, top_p, se
         messages.append({"role": "user", "content": message})
         
         response = ""
-        for message in client.chat_completion(
-            messages,
-            max_tokens=max_tokens,
-            stream=True,
-            temperature=temperature,
-            top_p=top_p,
-        ):
-            if hasattr(message.choices[0], 'delta'):
-                token = message.choices[0].delta.content
-            else:
-                token = message.choices[0].text  # 일부 모델은 'text' 속성을 사용할 수 있습니다.
-            
-            if token is not None:
-                response += token
-                yield response
-            
-        if not response:  # 응답이 비어있는 경우
+        for token in get_model_response(client, messages, max_tokens, temperature, top_p, selected_model):
+            response += token
+            yield response
+        
+        if not response:
             yield "모델이 응답을 생성하지 못했습니다. 다른 입력이나 모델을 시도해보세요."
     except Exception as e:
         error_msg = f"오류 발생: {str(e)}\n\n상세 오류:\n{traceback.format_exc()}"
         yield error_msg
 
+# 원래의 모델 목록 복원
 models = {
     "deepseek-ai/DeepSeek-Coder-V2-Instruct": "DeepSeek-Coder-V2-Instruct",
     "CohereForAI/c4ai-command-r-plus": "Cohere Command-R Plus",