import gradio as gr
from transformers import pipeline

# 모델 로드
try:
    generator = pipeline("text-generation", model="bigscience/bloomz-560m")  # 필요에 따라 모델 변경
except Exception as e:
    generator = None
    error_message = f"모델 로드 오류: {str(e)}"

# 응답 생성 함수
def generate_reply(review):
    if not generator:
        return f"모델을 로드할 수 없습니다. 오류: {error_message}"

    # 프롬프트 생성
    prompt = f"""
    다음은 고객 리뷰입니다:
    리뷰: "{review}"
    이 리뷰에 대해 정중하고 감사의 뜻을 담은 고객 서비스 팀의 답변을 작성하세요.
    답변 예: "고객님의 소중한 의견 감사합니다. 앞으로도 더 나은 서비스를 제공하기 위해 노력하겠습니다."
    """

    try:
        # 모델 호출
        result = generator(prompt, max_new_tokens=50, do_sample=True, temperature=0.7)

        # 디버깅 정보 출력
        debug_info = f"입력된 프롬프트: {prompt}\n모델 응답: {result}"

        if result and "generated_text" in result[0]:
            generated_text = result[0]["generated_text"]

            # 프롬프트 내용 제거
            if prompt in generated_text:
                generated_text = generated_text.replace(prompt, "").strip()

            return f"생성된 답변: {generated_text}\n\n[디버깅 정보]\n{debug_info}"
        else:
            return f"응답 처리 중 문제가 발생했습니다.\n\n[디버깅 정보]\n{debug_info}"

    except Exception as e:
        return f"API 호출 중 오류가 발생했습니다: {str(e)}"

# Gradio 인터페이스 설정
iface = gr.Interface(
    fn=generate_reply,
    inputs="text",
    outputs="text",
    title="Review Reply Generator",
    description="고객 리뷰를 입력하면 정중한 답변을 생성합니다."
)

# Space 실행
iface.launch()