Upload r_story_test.py

r_story_test.py

@@ -0,0 +1,133 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# ## 1. tsv full data load
+import pandas as pd
+
+df = pd.read_csv("sl_webtoon_full_data_sequential.tsv", sep="\t")
+
+print(df.head())
+print("전체 문장 수:", len(df))
+print("컬럼 목록:", df.columns.tolist())
+
+df['row_id'] = df.index  # 인덱스 컬럼 추가
+df['text'] = df.apply(
+    lambda x: f"[{x['에피소드']}] #{x['row_id']} {x['type']} {x['scene_text']}",
+    axis=1
+)
+texts = df['text'].tolist()
+print("최종 문장 수:", len(texts))
+
+# ## 2. RAG 문장 생성
+print("예시 5개:")
+for t in df['text'].head(5).tolist():
+    print("-", t)
+
+# ## 3. 한국어 임베딩 모델 로드, 벡터 db
+from langchain.vectorstores import FAISS
+from langchain.embeddings import HuggingFaceEmbeddings
+
+embedding_model = HuggingFaceEmbeddings(model_name='jhgan/ko-sroberta-multitask')
+
+db = FAISS.from_texts(texts, embedding_model)
+print(" 벡터DB 생성 완료. 총 문장 수:", len(texts))
+db.save_local("solo_leveling_faiss_ko")
+
+db = FAISS.load_local("solo_leveling_faiss_ko", embedding_model, allow_dangerous_deserialization=True)
+
+# 검색 테스트
+query = "마나석이 뭐지?"
+docs = db.similarity_search(query, k=5)
+for i, doc in enumerate(docs, 1):
+    print(f"[{i}] {doc.page_content}")
+
+# ## 4. LLM 로드 (CPU 전용)
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+from langchain.chains import RetrievalQA
+from langchain.prompts import PromptTemplate
+from langchain_community.llms import HuggingFacePipeline
+import torch
+
+# CPU로 강제
+generator = pipeline(
+    "text-generation",
+    model="kakaocorp/kanana-nano-2.1b-instruct",
+    device=-1   # ✅ CPU 사용
+)
+
+embedding_model = HuggingFaceEmbeddings(model_name='jhgan/ko-sroberta-multitask')
+vectorstore = FAISS.load_local("solo_leveling_faiss_ko", embedding_model, allow_dangerous_deserialization=True)
+
+model_name = "kakaocorp/kanana-nano-2.1b-instruct"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.float32   # ✅ CPU에서는 float32
+).to("cpu")                      # ✅ CPU 사용
+
+llm_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=128)
+llm = HuggingFacePipeline(pipeline=llm_pipeline)
+
+custom_prompt = PromptTemplate(
+    input_variables=["context", "question"],
+    template="다음 문맥을 참고하여 질문에 답하세요.\n\n문맥:\n{context}\n\n질문:\n{question}\n\n답변:"
+)
+
+qa_chain = RetrievalQA.from_chain_type(
+    llm=llm,
+    retriever=vectorstore.as_retriever(search_kwargs={"k": 5}),
+    chain_type="stuff",
+    return_source_documents=True,
+    chain_type_kwargs={"prompt": custom_prompt}
+)
+
+# 질문 테스트
+query = "성진우는 몇 급 헌터지?"
+result = qa_chain({"query": query})
+print("답변:", result["result"])
+print("\n참조 문서:")
+for doc in result["source_documents"]:
+    print(doc.page_content)
+
+# ## 5. 황동석 에피소드
+choices = [
+    "1: 황동석 무리를 모두 처치한다.",
+    "2: 진호를 포함한 황동석 무리를 모두 처치한다.",
+    "3: 전부 기절 시키고 살려둔다.",
+    "4: 시스템을 거부하고 그냥 도망친다."
+]
+
+print("\n[선택지]")
+for idx, choice in enumerate(choices, start=1):
+    print(f"{idx}. {choice}")
+
+user_idx = int(input("\n선택 번호 입력: ")) - 1
+user_choice = choices[user_idx]
+print(f"\n[사용자 선택]: {user_choice}")
+
+result = qa_chain({"query": user_choice})
+retrieved_context = "\n".join([doc.page_content for doc in result["source_documents"]])
+
+print("\n[검색된 근거 문서 예시]")
+print(retrieved_context[:600], "...") 
+
+prompt = f"""
+당신은 웹툰 '나 혼자만 레벨업'의 성진우입니다.
+현재 상황:
+{retrieved_context}
+사용자 선택: {user_choice}
+성진우의 말투로 간결하고 자연스러운 대사를 1~2문장 생성하세요.
+중복된 내용이나 비슷한 문장은 만들지 마세요.
+"""
+
+response = generator(
+    prompt, 
+    max_new_tokens=200, 
+    do_sample=True, 
+    temperature=0.6,
+    top_p=0.9,
+    return_full_text=False
+)[0]["generated_text"]
+
+print("\n[성진우 응답]")
+print(response)