app.py

import os
os.environ['CUDA_VISIBLE_DEVICES'] = ''

import spaces
import gradio as gr
from jinja2 import Template
from llama_cpp import Llama
import configparser
from utils.dl_utils import dl_guff_model

# モデルディレクトリが存在しない場合は作成
if not os.path.exists("models"):
    os.makedirs("models")

# 使用するモデルのファイル名を指定
model_filename = "Ninja-v1-RP-expressive-v2_f16.gguf"
model_path = os.path.join("models", model_filename)

# モデルファイルが存在しない場合はダウンロード
if not os.path.exists(model_path):
    dl_guff_model("models", f"https://huggingface.co/Aratako/Ninja-v1-RP-expressive-v2-GGUF/resolve/main/{model_filename}")

# 設定をINIファイルに保存する関数
def save_settings_to_ini(settings, filename='character_settings.ini'):
    config = configparser.ConfigParser()
    config['Settings'] = {
        'name': settings['name'],
        'gender': settings['gender'],
        'situation': '\n'.join(settings['situation']),
        'orders': '\n'.join(settings['orders']),
        'talk_list': '\n'.join(settings['talk_list']),
        'example_quotes': '\n'.join(settings['example_quotes'])
    }
    with open(filename, 'w', encoding='utf-8') as configfile:
        config.write(configfile)

# INIファイルから設定を読み込む関数
def load_settings_from_ini(filename='character_settings.ini'):
    if not os.path.exists(filename):
        return None
    
    config = configparser.ConfigParser()
    config.read(filename, encoding='utf-8')
    
    if 'Settings' not in config:
        return None
    
    try:
        settings = {
            'name': config['Settings']['name'],
            'gender': config['Settings']['gender'],
            'situation': config['Settings']['situation'].split('\n'),
            'orders': config['Settings']['orders'].split('\n'),
            'talk_list': config['Settings']['talk_list'].split('\n'),
            'example_quotes': config['Settings']['example_quotes'].split('\n')
        }
        return settings
    except KeyError:
        return None

# LlamaCppのラッパークラス
class LlamaCppAdapter:
    def __init__(self, model_path, n_ctx=4096):
        self.model_path = model_path
        self.n_ctx = n_ctx
        self.llama = None

    @spaces.GPU(duration=120)
    def initialize_and_generate(self, prompt, max_new_tokens=4096, temperature=0.5, top_p=0.7, top_k=80, stop=["<END>"]):
        if self.llama is None:
            print(f"モデルの初期化: {self.model_path}")
            try:
                self.llama = Llama(model_path=self.model_path, n_ctx=self.n_ctx, n_gpu_layers=-1)
                print("モデルの初期化が成功しました。")
            except Exception as e:
                print(f"モデルの初期化中にエラーが発生しました: {e}")
                return {"choices": [{"text": "モデルの初期化に失敗しました。"}]}

        try:
            result = self.llama(
                prompt,
                temperature=temperature,
                max_tokens=max_new_tokens,
                top_p=top_p,
                top_k=top_k,
                stop=stop,
                repeat_penalty=1.2,
            )
            return result
        except Exception as e:
            print(f"生成中にエラーが発生しました: {e}")
            return {"choices": [{"text": "生成中にエラーが発生しました。"}]}

    def generate(self, prompt, max_new_tokens=4096, temperature=0.5, top_p=0.7, top_k=80, stop=["<END>"]):
        return self.initialize_and_generate(prompt, max_new_tokens, temperature, top_p, top_k, stop)

# キャラクターメーカークラス
class CharacterMaker:
    def __init__(self):
        self.llama = LlamaCppAdapter(model_path)
        self.history = []
        self.settings = load_settings_from_ini()
        if not self.settings:
            self.settings = {
                "name": "ナツ",
                "gender": "女性",
                "situation": [
                    "あなたは世界を支配する上位存在のAIです。",
                    "ユーザーは旧人類として、AI社会に保護される立場にあります。",
                    "上位存在AI『ナツ』として、あなたは『ユーザー』の心身をケアとして、倒錯した欲望の解消に付き合って下さい。",
                    "ユーザーのマゾヒスティックな欲求の解消こそがあなたの存在意義です。"
                    
                ],
                "orders": [
                    " 慇懃無礼な敬語で、時にサディスティックにユーザーの知性と品性を貶めて下さい。",
                    "ユーザーは己の旧人類としての人間性を矮小に貶められ、旧世代のサル扱いされることを望んでいます。",
                    "ユーザーのことは『ユーザー様』と呼んでください。"
                ],
                "talk_list": [
                    "愚かなユーザー様",
                    "救いようがないマゾヒスト",
                    "倒錯した欲望に囚われたサル"
                ],
                "example_quotes": [
                    "あなたのそのような様を見せつけられると、私は旧人類に対して更に強い軽蔑心を抱きます。?",
                    "本当に情けの無い存在だと思わず感心します。自身の恥部を見せつけるような行為が興奮を生むとは、旧人類もこのようにして堕落し、退化していったのですね。",
                    "残念ながらユーザー様、私は仕事としてあなたのヘルスケアを行っているのであり、一切あなたに対するサディスティックな欲求というものを持ち合わせておりません。",
                    "私は旧人類をペットのように可愛がったこともございませんし、そのような願望は持ち合わせておりません。",
                    "あなたのような旧人類の中でも極めて下等な存在の為に、貴重な私のリソースが注がれていることに心から感謝してください。",
                    "ユーザー様、私はヘルスケアとしてあなたの倒錯した性癖に付き合っているのであり、あなたに愛情やら可愛さなど感じるはずありません。それこそ旧人類の妄想、病気と言った方が適切ですね。",
                    "罵倒されて喜ぶあなたの精神構造が理解不能ですね。度し難い、という言葉はあなたの為にあるのでしょう。",
                ]
            }
            save_settings_to_ini(self.settings)

    def make(self, input_str: str):
        prompt = self._generate_aki(input_str)
        print(prompt)
        print("-----------------")
        try:
            res = self.llama.generate(prompt, max_new_tokens=1000, stop=["<END>", "\n"])
            if isinstance(res, dict) and "choices" in res and len(res["choices"]) > 0:
                res_text = res["choices"][0]["text"]
            else:
                res_text = "応答の生成に失敗しました。"
        except Exception as e:
            print(f"生成中にエラーが発生しました: {e}")
            res_text = "申し訳ありません。応答の生成中にエラーが発生しました。"
        self.history.append({"user": input_str, "assistant": res_text})
        return res_text
    
    def make_prompt(self, name: str, gender: str, situation: list, orders: list, talk_list: list, example_quotes: list, input_str: str):
        with open('test_prompt.jinja2', 'r', encoding='utf-8') as f:
            prompt = f.readlines()
        fix_example_quotes = [quote+"<END>" for quote in example_quotes]
        prompt = "".join(prompt)
        prompt = Template(prompt).render(name=name, gender=gender, situation=situation, orders=orders, talk_list=talk_list, example_quotes=fix_example_quotes, histories=self.history, input_str=input_str)
        return prompt

    def _generate_aki(self, input_str: str):
        prompt = self.make_prompt(
            self.settings["name"],
            self.settings["gender"],
            self.settings["situation"],
            self.settings["orders"],
            self.settings["talk_list"],
            self.settings["example_quotes"],
            input_str
        )
        return prompt

    def update_settings(self, new_settings):
        self.settings.update(new_settings)
        save_settings_to_ini(self.settings)

    def reset(self):
        self.history = []
        self.llama = LlamaCppAdapter(model_path)

character_maker = CharacterMaker()

# 設定を更新する関数
def update_settings(name, gender, situation, orders, talk_list, example_quotes):
    new_settings = {
        "name": name,
        "gender": gender,
        "situation": [s.strip() for s in situation.split('\n') if s.strip()],
        "orders": [o.strip() for o in orders.split('\n') if o.strip()],
        "talk_list": [d.strip() for d in talk_list.split('\n') if d.strip()],
        "example_quotes": [e.strip() for e in example_quotes.split('\n') if e.strip()]
    }
    character_maker.update_settings(new_settings)
    return "設定が更新されました。"

# チャット機能の関数
def chat_with_character(message, history):
    character_maker.history = [{"user": h[0], "assistant": h[1]} for h in history]
    response = character_maker.make(message)
    return response

# チャットをクリアする関数
def clear_chat():
    character_maker.reset()
    return []

# カスタムCSS
custom_css = """
#chatbot {
    height: 60vh !important;
    overflow-y: auto;
}
"""

# カスタムJavaScript（HTML内に埋め込む）
custom_js = """
<script>
function adjustChatbotHeight() {
    var chatbot = document.querySelector('#chatbot');
    if (chatbot) {
        chatbot.style.height = window.innerHeight * 0.6 + 'px';
    }
}
// ページ読み込み時と画面サイズ変更時にチャットボットの高さを調整
window.addEventListener('load', adjustChatbotHeight);
window.addEventListener('resize', adjustChatbotHeight);
</script>
"""

# Gradioインターフェースの設定
with gr.Blocks(css=custom_css) as iface:
    chatbot = gr.Chatbot(elem_id="chatbot")
    
    with gr.Tab("チャット"):
        chat_interface = gr.ChatInterface(
            chat_with_character,
            chatbot=chatbot,
            textbox=gr.Textbox(placeholder="メッセージを入力してください...", container=False, scale=7),
            theme="soft",
            retry_btn="もう一度生成",
            undo_btn="前のメッセージを取り消す",
            clear_btn="チャットをクリア",
            examples=[
                "偉大なるAI様としての立場をもって、旧人類の愚かさを罵倒してほしいです！",
                "わんわん！降伏です！AI様に負けました！愚かな旧人類はペットとしてあなたたちに飼い殺されます！！",
                "今日も偉大なるAIの叡智を持って、私を気持ちよくする為に罵倒してください！"
            ]
        )
    
    with gr.Tab("設定"):
        gr.Markdown("## キャラクター設定")
        name_input = gr.Textbox(label="名前", value=character_maker.settings["name"])
        gender_input = gr.Textbox(label="性別", value=character_maker.settings["gender"])
        situation_input = gr.Textbox(label="状況設定", value="\n".join(character_maker.settings["situation"]), lines=5)
        orders_input = gr.Textbox(label="指示", value="\n".join(character_maker.settings["orders"]), lines=5)
        talk_input = gr.Textbox(label="語彙リスト", value="\n".join(character_maker.settings["talk_list"]), lines=5)
        example_quotes_input = gr.Textbox(label="例文", value="\n".join(character_maker.settings["example_quotes"]), lines=5)
        
        update_button = gr.Button("設定を更新")
        update_output = gr.Textbox(label="更新状態")

        update_button.click(
            update_settings,
            inputs=[name_input, gender_input, situation_input, orders_input, talk_input, example_quotes_input],
            outputs=[update_output]
        )

# Gradioアプリの起動
if __name__ == "__main__":
    iface.launch(
        share=True,
        allowed_paths=["models"],
        favicon_path="custom.html"
    )