Update app.py

app.py

@@ -1,7 +1,13 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+from transformers import pipeline
+import os
+from typing import List, Tuple, Generator
 import concurrent.futures
 
+# Hugging Face 토큰 설정
+os.environ["TOKENIZERS_PARALLELISM"] = "false"  # 경고 메시지 방지
+HF_TOKEN = os.getenv("HF_TOKEN")
+
 # Available LLM models
 LLM_MODELS = {
     "Llama-3.3": "meta-llama/Llama-3.3-70B-Instruct",
@@ -20,24 +26,27 @@ DEFAULT_MODELS = [
     "mistralai/Mistral-Nemo-Instruct-2407"
 ]
 
-clients = {model: InferenceClient(model) for model in LLM_MODELS.values()}
+# Pipeline 초기화
+pipes = {}
+for model_name in LLM_MODELS.values():
+    try:
+        pipes[model_name] = pipeline(
+            "text-generation",
+            model=model_name,
+            token=HF_TOKEN,
+            device_map="auto"
+        )
+    except Exception as e:
+        print(f"Failed to load model {model_name}: {str(e)}")
 
-def process_file(file):
+def process_file(file) -> str:
     if file is None:
         return ""
     if file.name.endswith(('.txt', '.md')):
         return file.read().decode('utf-8')
     return f"Uploaded file: {file.name}"
 
-def respond_single(
-    client,
-    message,
-    history,
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
+def format_messages(message: str, history: List[Tuple[str, str]], system_message: str) -> List[dict]:
     messages = [{"role": "system", "content": system_message}]
     
     for user, assistant in history:
@@ -47,34 +56,50 @@ def respond_single(
             messages.append({"role": "assistant", "content": assistant})
     
     messages.append({"role": "user", "content": message})
-    
-    response = ""
+    return messages
+
+def generate_response(
+    pipe,
+    messages: List[dict],
+    max_tokens: int,
+    temperature: float,
+    top_p: float
+) -> Generator[str, None, None]:
     try:
-        for msg in client.chat_completion(
-            messages,
-            max_tokens=max_tokens,
-            stream=True,
+        formatted_prompt = "\n".join([f"{msg['role']}: {msg['content']}" for msg in messages])
+        
+        response = pipe(
+            formatted_prompt,
+            max_new_tokens=max_tokens,
             temperature=temperature,
             top_p=top_p,
-        ):
-            token = msg.choices[0].delta.content
-            response += token
-            yield response
+            do_sample=True,
+            pad_token_id=50256,
+            num_return_sequences=1,
+            streaming=True
+        )
+        
+        generated_text = ""
+        for output in response:
+            new_text = output[0]['generated_text'][len(formatted_prompt):].strip()
+            generated_text = new_text
+            yield generated_text
+            
     except Exception as e:
         yield f"Error: {str(e)}"
 
 def respond_all(
-    message,
+    message: str,
     file,
-    history1,
-    history2,
-    history3,
-    selected_models,
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
+    history1: List[Tuple[str, str]],
+    history2: List[Tuple[str, str]],
+    history3: List[Tuple[str, str]],
+    selected_models: List[str],
+    system_message: str,
+    max_tokens: int,
+    temperature: float,
+    top_p: float,
+) -> Tuple[Generator[str, None, None], Generator[str, None, None], Generator[str, None, None]]:
     if file:
         file_content = process_file(file)
         message = f"{message}\n\nFile content:\n{file_content}"
@@ -82,21 +107,14 @@ def respond_all(
     while len(selected_models) < 3:
         selected_models.append(selected_models[-1])
 
-    def generate(client, history):
-        return respond_single(
-            client,
-            message,
-            history,
-            system_message,
-            max_tokens,
-            temperature,
-            top_p,
-        )
+    def generate(pipe, history):
+        messages = format_messages(message, history, system_message)
+        return generate_response(pipe, messages, max_tokens, temperature, top_p)
 
     return (
-        generate(clients[selected_models[0]], history1),
-        generate(clients[selected_models[1]], history2),
-        generate(clients[selected_models[2]], history3),
+        generate(pipes[selected_models[0]], history1),
+        generate(pipes[selected_models[1]], history2),
+        generate(pipes[selected_models[2]], history3),
     )
 
 with gr.Blocks() as demo:
@@ -186,4 +204,7 @@ with gr.Blocks() as demo:
     )
 
 if __name__ == "__main__":
+    # Hugging Face 토큰이 설정되어 있는지 확인
+    if not HF_TOKEN:
+        print("Warning: HF_TOKEN environment variable is not set")
     demo.launch()