Update app.py

app.py

@@ -4,6 +4,8 @@ import openai
 import anthropic
 import os
 from typing import Optional
+import transformers
+import torch
 
 #############################
 # [기본코드] - 수정/삭제 불가
@@ -171,25 +173,27 @@ def respond_claude_qna(
         return f"예상치 못한 오류가 발생했습니다: {str(e)}"
 
 #############################
-# [추가코드] - Llama-3.3-70B-Instruct / Llama-3.2-3B-Instruct 적용
+# [추가코드] - Llama-3.3-70B-Instruct / Llama-3.2-3B-Instruct 적용 (transformers.pipeline 방식)
 #############################
 
 def get_llama_client(model_choice: str):
     """
-    선택된 Llama 모델에 맞춰 InferenceClient 생성.
-    토큰은 환경 변수에서 가져옴.
+    선택된 Llama 모델에 맞춰 transformers의 text-generation 파이프라인을 생성.
     """
-    hf_token = os.getenv("HF_TOKEN")
-    if not hf_token:
-        raise ValueError("HuggingFace API 토큰이 필요합니다.")
-    
     if model_choice == "Llama-3.3-70B-Instruct":
-        model_id = "Llama-3.3-70B-Instruct"
+        model_id = "meta-llama/Llama-3.3-70B-Instruct"
     elif model_choice == "Llama-3.2-3B-Instruct":
-        model_id = "Llama-3.2-3B-Instruct"
+        model_id = "meta-llama/Llama-3.2-3B-Instruct"
     else:
         raise ValueError("유효하지 않은 모델 선택입니다.")
-    return InferenceClient(model_id, token=hf_token)
+    
+    pipeline_llama = transformers.pipeline(
+        "text-generation",
+        model=model_id,
+        model_kwargs={"torch_dtype": torch.bfloat16},
+        device_map="auto",
+    )
+    return pipeline_llama
 
 def respond_llama_qna(
     question: str,
@@ -200,27 +204,27 @@ def respond_llama_qna(
     model_choice: str
 ):
     """
-    선택된 Llama 모델을 이용해 한 번의 질문(question)에 대한 답변을 반환하는 함수.
+    선택된 Llama 모델을 이용해 한 번의 질문(question)에 대한 답변을 transformers 파이프라인으로 반환하는 함수.
+    system_message와 question을 하나의 프롬프트로 결합하여 생성합니다.
     """
     try:
-        client = get_llama_client(model_choice)
+        pipeline_llama = get_llama_client(model_choice)
     except ValueError as e:
         return f"오류: {str(e)}"
     
-    messages = [
-        {"role": "system", "content": system_message},
-        {"role": "user", "content": question}
-    ]
+    # system_message와 question을 연결하여 프롬프트 생성
+    prompt = system_message.strip() + "\n" + question.strip()
     
     try:
-        response_full = client.chat_completion(
-            messages,
-            max_tokens=max_tokens,
+        outputs = pipeline_llama(
+            prompt,
+            max_new_tokens=max_tokens,
             temperature=temperature,
             top_p=top_p,
         )
-        assistant_message = response_full.choices[0].message.content
-        return assistant_message
+        # 생성된 텍스트를 추출 (전체 프롬프트 이후의 텍스트만 반환할 수도 있음)
+        generated_text = outputs[0]["generated_text"]
+        return generated_text
     except Exception as e:
         return f"오류가 발생했습니다: {str(e)}"