model test

app/main.py

@@ -1,181 +1,95 @@
-from fastapi import FastAPI, Request
-from pydantic import BaseModel
-from typing import List
-from transformers import pipeline
-from PIL import Image
-import re, os
+from fastapi import FastAPI
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
 import torch
 
 app = FastAPI()
 
-# 1. LLM 파이프라인 초기화 (SmolLM3 모델)
 print("[torch] is available:", torch.cuda.is_available())
 print("[device] default:", torch.device("cuda" if torch.cuda.is_available() else "cpu"))
-llm = pipeline("text-generation", model="HuggingFaceTB/SmolLM3-3B", device=0 if torch.cuda.is_available() else -1)
 
-# 2. 감정 및 상황별 이미지 매핑
-'''
-이미지 매핑 예시
------------------------------
-이미지를 다음과 같이 집어넣으면 (./asset/face/)
-happy.png
-sad.png
-angry.png
-=>
-이런 딕셔너리 형식으로 반환됨
-{
-    "happy": "happy.png",
-    "sad": "sad.png",
-    "angry": "angry.png"
-}
-=>
-모델 출력의 감정 부분에 대응되는 이미지 출력
-'''
-def load_faces(face_dir="assets/face"):
-    if not os.path.exists(face_dir):
-        os.makedirs(face_dir)
-    emotion_to_face = {}
-    for filename in os.listdir(face_dir):
-        if filename.endswith(".png"):
-            emotion = os.path.splitext(filename)[0]  # 'happy.png' → 'happy'
-            emotion_to_face[emotion] = filename      # "happy": "happy.png"
-    return emotion_to_face
-
-def load_bgs(bg_dir="assets/bg"):
-    if not os.path.exists(bg_dir):
-        os.makedirs(bg_dir)
-    situation_to_bg = {}
-    for filename in os.listdir(bg_dir):
-        if filename.endswith(".png"):
-            emotion = os.path.splitext(filename)[0]  # 'happy.png' → 'happy'
-            situation_to_bg[emotion] = filename      # "happy": "happy.png"
-    return situation_to_bg
-
-emotion_to_face = load_faces()
-situation_to_bg = load_bgs()
-
-# 3. 출력 라인 파싱 함수
-def parse_output(text: str):
-    pattern = r'"(.*?)"\s*\(emotion:\s*(\w+),\s*situation:\s*(\w+)\)'
-    results = []
-    for line in text.strip().split('\n'):
-        match = re.match(pattern, line.strip())
-        if match:
-            results.append({
-                "text": match.group(1),
-                "emotion": match.group(2),
-                "situation": match.group(3)
-            })
-    return results
-
-# 4. 이미지 합성 함수
-def combine_images(bg_path, face_path):
-    try:
-        bg = Image.open(bg_path).convert("RGBA")
-    except FileNotFoundError:
-        print(f"[warning] 배경 이미지 없음: {bg_path}")
-        return None
-    try:
-        face = Image.open(face_path).convert("RGBA")
-    except FileNotFoundError:
-        print(f"[warning] 캐릭터 이미지 없음: {face_path}")
-        return None
-    # 이미지 합성
-    bg.paste(face, (0, 0), face)
-    return bg
+# 모델 로드
+# https://huggingface.co/EleutherAI/polyglot-ko-1.3b
+model_id = "EleutherAI/polyglot-ko-1.3b"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id)
+llm = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    device=0
+)
+
+# 챗봇 프롬프트 생성
+chat_history = []
 
-# 5. 챗봇 처리 함수 (Gradio용)
-'''
-지금까지 대화 내용을 모두 프롬프트로 넣어서 대화내용을 기억하도록 함
-'''
-def build_prompt(chat_history, user_msg):
-    system_prompt = (
-        "You are Aria, a cheerful and expressive fantasy mage."
-        " Respond in multiple steps if needed."
-        " Format: \"text\" (emotion: tag, situation: tag)"
+def build_prompt(history, user_msg):
+    prompt = (
+            "[시작]\n"
+            "당신은 마법사 아리아(Aria)입니다.\n"
+            "규칙:\n"
+            "- 항상 한 문장만 말합니다.\n"
+            "- 사용자 발화를 반복하거나 따라하지 않습니다.\n"
+            "- 영어, 인용문, 중괄호, 특수기호를 사용하지 않습니다.\n"
+            "- 사용자 질문에만 반응하고 혼잣말을 하지 않습니다.\n"
+            "- 항상 한국어만 사용해서 대답합니다.\n"
+            "대화 예시:\n"
+            "User: 안녕!\n"
+            "Aria: 안녕하세요, 무엇을 도와드릴까요?\n"
+            "User: 이름이 뭐야?\n"
+            "Aria: 저는 아리아라고 해요."
+            )
+    for turn in history[-2:]:  # 최근 2턴만 사용
+        if turn["role"] == "user":
+            prompt += turn['text']
+        else:
+            prompt += turn['text']
+    prompt += user_msg
+    return prompt
+
+def character_chat(user_msg):
+    prompt = build_prompt(chat_history, user_msg)
+    outputs = llm(
+        prompt,
+        do_sample=True, 
+        max_new_tokens=20,
+        temperature=0.7,
+        top_p=0.8,
+        repetition_penalty=1.5,
+        eos_token_id=tokenizer.eos_token_id,
+        return_full_text=False
     )
+    response = outputs[0]['generated_text'].strip()
+    return response
 
-    dialogue = ""
-    for item in chat_history:
-        if item["role"] == "user":
-            dialogue += f"User: {item['text']}\n"
-        elif item["role"] == "bot":
-            dialogue += f"Aria: {item['text']}\n"
-
-    dialogue += f"User: {user_msg}\nAria:"
-    return system_prompt + "\n" + dialogue
-
-def character_chat(prompt):
-    full_prompt = build_prompt(chat_history, prompt)
-
-    #raw_output = llm(full_prompt, max_new_tokens=300)[0]['generated_text']
-    raw_output = '"우오아" (emotion: tag, situation: tag)'
-    parsed = parse_output(raw_output)
-
-    result_outputs = []
-    for i, item in enumerate(parsed):
-        face = emotion_to_face.get(item['emotion'], "neutral.png")
-        bg = situation_to_bg.get(item['situation'], "default.jpg")
-        composite = combine_images(os.path.join("assets/bg", bg), os.path.join("assets/face", face))
-        img_path = None #이미지가 없으면 출력 안함
-        if composite:
-            img_path = f"static/output_{i}.png"
-            composite.save(img_path)
-        result_outputs.append((item['text'], img_path))
-
-    return result_outputs
-
-# 6. Gradio UI with chat history
-chat_history = []
-
+# Gradio 인터페이스
 with gr.Blocks(css="""
 .chat-box { max-height: 500px; overflow-y: auto; padding: 10px; border: 1px solid #ccc; border-radius: 10px; }
 .bubble-left { background-color: #f1f0f0; border-radius: 10px; padding: 10px; margin: 5px; max-width: 70%; float: left; clear: both; }
 .bubble-right { background-color: #d1e7ff; border-radius: 10px; padding: 10px; margin: 5px; max-width: 70%; float: right; clear: both; text-align: right; }
-.image-preview { margin: 5px 0; max-width: 100%; border-radius: 10px; }
 """) as demo:
-    gr.Markdown("챗봇")
+    gr.Markdown("### 아리아와 대화하기")
     with gr.Column():
-        chat_output = gr.HTML(value="<div class='chat-box' id='chat-box'></div>")
-        user_input = gr.Textbox(label="Say something to Aria", placeholder="Type here and press Enter")
+        chat_output = gr.HTML(elem_id="chat-box")
+        user_input = gr.Textbox(label="메시지 입력", placeholder="Aria에게 말을 걸어보세요")
 
     def render_chat():
         html = ""
         for item in chat_history:
-            if item['role'] == 'user':
+            if item["role"] == "user":
                 html += f"<div class='bubble-right'>{item['text']}</div>"
-            elif item['role'] == 'bot':
-                bubble = f"<div class='bubble-left'>{item['text']}"
-                if 'image' in item and item['image']:
-                    bubble += f"<br><img class='image-preview' src='{item['image']}'>"
-                bubble += "</div>"
-                html += bubble
-        return html
+            elif item["role"] == "bot":
+                html += f"<div class='bubble-left'>{item['text']}</div>"
+        return gr.update(value=html)
 
     def on_submit(user_msg):
         chat_history.append({"role": "user", "text": user_msg})
+        yield render_chat(), ""
+        response = character_chat(user_msg)
+        chat_history.append({"role": "bot", "text": response})
+        yield render_chat(), ""
 
-        bot_results = character_chat(user_msg)
-
-        for item in bot_results:
-            try:
-                text, image_path = item  # unpack 시도
-            except (ValueError, TypeError):
-                # unpack 안되면 기본값 처리 (이미지 없이)
-                text = str(item)
-                image_path = None
-
-            chat_entry = {"role": "bot", "text": text}
-            if image_path:
-                chat_entry["image"] = image_path
-
-            chat_history.append(chat_entry)
-
-        new_chat_html = render_chat()
-        return f"<div class='chat-box' id='chat-box'>{new_chat_html}</div>", ""
-
-    user_input.submit(on_submit, inputs=user_input, outputs=[chat_output, user_input])
+    user_input.submit(on_submit, inputs=user_input, outputs=[chat_output, user_input], queue=True)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()