app.py

import pandas as pd
import torch
from flask import Flask, request, Response
from transformers import AutoTokenizer, AutoModelForCausalLM
from dicttoxml import dicttoxml
import traceback
import re
from threading import Lock

app = Flask(__name__)

# --- 1. 디바이스 설정 ---
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"사용 디바이스: {device}")

torch.set_grad_enabled(False)

# --- 2. 모델 로드 ---
print("토크나이저 로딩 중...")
tokenizer = AutoTokenizer.from_pretrained(
    "LiquidAI/LFM2.5-1.2B-Instruct",
    trust_remote_code=True
)

print("모델 로딩 중...")
try:
    # 8bit 로드 시도
    model = AutoModelForCausalLM.from_pretrained(
        "LiquidAI/LFM2.5-1.2B-Instruct",
        device_map="auto",
        load_in_8bit=True,
        trust_remote_code=True
    )
    print("8bit 로딩 성공")
except:
    # 실패 시 일반 로드
    model = AutoModelForCausalLM.from_pretrained(
        "LiquidAI/LFM2.5-1.2B-Instruct",
        trust_remote_code=True
    ).to(device)
    print("일반 로딩 사용")

# torch 2.0 이상이면 컴파일
try:
    model = torch.compile(model)
    print("torch.compile 적용 완료")
except:
    print("torch.compile 미적용 (지원 안함)")

print("모델 로딩 완료!")

# --- 3. 데이터셋 로드 ---
try:
    df = pd.read_excel('dataset.xlsx')
    knowledge_list = df['데이터셋에 넣을 내용(*)'].tolist()
except Exception as e:
    print(f"데이터셋 로드 에러: {e}")
    knowledge_list = []

# --- 4. 동시 요청 제한용 Lock (구조 유지) ---
request_lock = Lock()

# --- 5. 질문과 관련된 지식 검색 (기존 방식 유지) ---
def find_relevant_context(query, top_n=2):
    query_words = query.replace(" ", "").lower()
    relevant_sentences = []
    for s in knowledge_list:
        s_text = str(s).replace(" ", "").replace("\n", "").lower()
        if any(word.replace(" ", "") in s_text for word in query.split()):
            relevant_sentences.append(s)
    return " ".join(str(s) for s in relevant_sentences[:top_n]) if relevant_sentences else ""

# --- 6. Sayknow 답변 생성 ---
def ask_sayknow(query):
    try:
        context = find_relevant_context(query)

        persona_guide = (
            "너는 지식 기반 한국어 챗봇 Sayknow야. 자기소개 질문에는 '저는 Sayknow입니다.'라고 답해. "
            "그 외에는 아래 참고해서 정확하고 자연스러운 한국어 문장으로 80자 이내로 답해.\n"
            "예시: Q: 분수의 덧셈이 뭐야?\nA: 분모가 같을 때 분자끼리 더하면 됩니다.\n"
        )

        info = context if context else "정보 없음"

        prompt = f"{persona_guide}---\n[정보]\n{info}\n[질문]\n{query}\n[답변] "

        tokenizer.pad_token = tokenizer.eos_token

        encoded_input = tokenizer.encode_plus(
            prompt,
            return_tensors='pt',
            truncation=True,
            padding=True
        )

        input_ids = encoded_input['input_ids'].to(device)
        attention_mask = encoded_input['attention_mask'].to(device)

        model.eval()

        gen_ids = model.generate(
            input_ids,
            attention_mask=attention_mask,
            max_new_tokens=60,      # 줄임
            min_length=5,
            repetition_penalty=1.2,
            do_sample=True,
            top_k=30,
            top_p=0.8,
            temperature=0.5,
            num_beams=1,
            pad_token_id=tokenizer.pad_token_id
        )

        raw_response = tokenizer.decode(gen_ids[0], skip_special_tokens=True)

        # --- 답변 추출 ---
        answer = raw_response.replace(prompt, '').strip()

        if "답변:" in answer:
            answer = answer.split("답변:", 1)[1].strip()

        # --- 후처리 (5번 유지 요청대로 그대로 유지) ---
        answer = re.sub(r"[^가-힣0-9 .,!?~\n]", "", answer)
        answer = re.sub(r"([.,!?~])\1{2,}", r"\1", answer)
        answer = re.sub(r"[a-zA-Z]+", "", answer)
        answer = re.sub(r"[=^*/\\]+", "", answer)
        answer = re.sub(r"\s+", " ", answer).strip()

        # 80자 제한
        answer = answer[:80]

        if answer and answer[-1] not in ".!?":
            answer += "."
        elif not answer:
            answer = "죄송합니다. 질문에 대한 답변을 찾을 수 없습니다."

        return answer

    except Exception as e:
        print(f"ask_sayknow 에러: {e}")
        traceback.print_exc()
        return f"내부 오류: {str(e)}"

# --- 7. API (XML 응답) ---
@app.route('/chatapi.html', methods=['GET'])
@app.route('/index.html', methods=['GET'])
def chat_api():
    query = request.args.get('askdata', '')

    if not query:
        result = {"status": "error", "message": "No data"}
        xml_output = dicttoxml(result, custom_root='SayknowAPI', attr_type=False)
        return Response(xml_output, mimetype='text/xml')

    # 6번 유지 요청 → Lock 전체 유지
    with request_lock:
        try:
            answer = ask_sayknow(query)
            result = {
                "service": "Sayknow",
                "question": query,
                "answer": answer
            }
        except Exception as e:
            print(f"chat_api 에러: {e}")
            traceback.print_exc()
            result = {
                "service": "Sayknow",
                "question": query,
                "answer": f"에러 발생: {str(e)}",
                "error": str(e)
            }

    xml_output = dicttoxml(result, custom_root='SayknowAPI', attr_type=False)
    return Response(xml_output, mimetype='text/xml')


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)