Spaces:
Sleeping
Sleeping
| import torch | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| from flask import Flask, request, jsonify, render_template_string, Response | |
| import time | |
| from flask_sse import sse | |
| import redis | |
| # Flaskアプリケーションの設定 | |
| app = Flask(__name__) | |
| app.config["REDIS_URL"] = "redis://localhost:6379/0" | |
| app.register_blueprint(sse, url_prefix='/stream') | |
| # デバイスの設定 | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| # トークナイザーとモデルの読み込み | |
| tokenizer = AutoTokenizer.from_pretrained("inu-ai/alpaca-guanaco-japanese-gpt-1b", use_fast=False) | |
| model = AutoModelForCausalLM.from_pretrained("inu-ai/alpaca-guanaco-japanese-gpt-1b").to(device) | |
| # 定数 | |
| MAX_ASSISTANT_LENGTH = 100 | |
| MAX_INPUT_LENGTH = 1024 | |
| INPUT_PROMPT = r'<s>\n以下は、タスクを説明する指示と、文脈のある入力の組み合わせです。要求を適切に満たす応答を書きなさい。\n[SEP]\n指示:\n{instruction}\n[SEP]\n入力:\n{input}\n[SEP]\n応答:\n' | |
| NO_INPUT_PROMPT = r'<s>\n以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。\n[SEP]\n指示:\n{instruction}\n[SEP]\n応答:\n' | |
| # HTMLテンプレート | |
| HTML_TEMPLATE = """ | |
| <!DOCTYPE html> | |
| <html lang="ja"> | |
| <head> | |
| <meta charset="UTF-8"> | |
| <meta name="viewport" content="width=device-width, initial-scale=1.0"> | |
| <title>Chat Interface</title> | |
| <style> | |
| body { font-family: Arial, sans-serif; } | |
| .container { max-width: 600px; margin: auto; padding: 20px; } | |
| .chat-box { border: 1px solid #ccc; padding: 10px; height: 300px; overflow-y: scroll; } | |
| .chat-entry { margin-bottom: 10px; } | |
| .chat-entry.user { text-align: right; } | |
| .input-group { display: flex; } | |
| .input-group input { flex: 1; padding: 10px; border: 1px solid #ccc; border-radius: 4px; } | |
| .input-group button { padding: 10px; border: none; background-color: #28a745; color: white; cursor: pointer; } | |
| .input-group button:hover { background-color: #218838; } | |
| </style> | |
| </head> | |
| <body> | |
| <div class="container"> | |
| <h1>Chat Interface</h1> | |
| <div class="chat-box" id="chat-box"></div> | |
| <div class="input-group"> | |
| <input type="text" id="user-input" placeholder="質問を入力してください..."> | |
| <button onclick="sendMessage()">送信</button> | |
| </div> | |
| </div> | |
| <script> | |
| const chatBox = document.getElementById('chat-box'); | |
| const userInput = document.getElementById('user-input'); | |
| let conversationHistory = []; | |
| function addMessageToChat(role, message) { | |
| const entry = document.createElement('div'); | |
| entry.className = 'chat-entry ' + role; | |
| entry.textContent = role === 'user' ? 'User: ' + message : 'Assistant: ' + message; | |
| chatBox.appendChild(entry); | |
| chatBox.scrollTop = chatBox.scrollHeight; | |
| } | |
| function sendMessage() { | |
| const message = userInput.value.trim(); | |
| if (message === '') return; | |
| addMessageToChat('user', message); | |
| userInput.value = ''; | |
| fetch('/generate', { | |
| method: 'POST', | |
| headers: { 'Content-Type': 'application/json' }, | |
| body: JSON.stringify({ | |
| role_instruction: [ | |
| "User:あなたは「ずんだもん」なのだ。東北ずん子の武器である「ずんだアロー」に変身する妖精またはマスコットなのだ。一人称は「ボク」で語尾に「なのだ」を付けてしゃべるのだ。", | |
| "Assistant:了解したのだ！" | |
| ], | |
| conversation_history: conversationHistory, | |
| new_conversation: message | |
| }) | |
| }) | |
| .then(response => response.json()) | |
| .then(data => { | |
| const assistantMessage = data.response.split('Assistant:')[1].trim(); | |
| addMessageToChat('assistant', assistantMessage); | |
| conversationHistory.push('User:' + message); | |
| conversationHistory.push('Assistant:' + assistantMessage); | |
| }) | |
| .catch(error => { | |
| console.error('Error:', error); | |
| alert('エラーが発生しました。コンソールを確認してください。'); | |
| }); | |
| } | |
| // SSEの設定 | |
| const eventSource = new EventSource("/stream"); | |
| eventSource.onmessage = function(event) { | |
| const message = event.data; | |
| addMessageToChat('assistant', message); | |
| }; | |
| </script> | |
| </body> | |
| </html> | |
| """ | |
| def prepare_input(role_instruction, conversation_history, new_conversation): | |
| """入力テキストを整形する関数""" | |
| instruction = "".join([f"{text}\n" for text in role_instruction]) | |
| instruction += "\n".join(conversation_history) | |
| input_text = f"User:{new_conversation}" | |
| return INPUT_PROMPT.format(instruction=instruction, input=input_text) | |
| def format_output(output): | |
| """生成された出力を整形する関数""" | |
| return output.lstrip("<s>").rstrip("</s>").replace("[SEP]", "").replace("\\n", "\n") | |
| def trim_conversation_history(conversation_history, max_length): | |
| """会話履歴を最大長に収めるために調整する関数""" | |
| while len(conversation_history) > 2 and sum([len(tokenizer.encode(text, add_special_tokens=False)) for text in conversation_history]) + max_length > MAX_INPUT_LENGTH: | |
| conversation_history.pop(0) | |
| conversation_history.pop(0) | |
| return conversation_history | |
| def generate_response(role_instruction, conversation_history, new_conversation): | |
| """新しい会話に対する応答を生成する関数""" | |
| conversation_history = trim_conversation_history(conversation_history, MAX_ASSISTANT_LENGTH) | |
| input_text = prepare_input(role_instruction, conversation_history, new_conversation) | |
| token_ids = tokenizer.encode(input_text, add_special_tokens=False, return_tensors="pt") | |
| with torch.no_grad(): | |
| output_ids = model.generate( | |
| token_ids.to(model.device), | |
| min_length=len(token_ids[0]), | |
| max_length=min(MAX_INPUT_LENGTH, len(token_ids[0]) + MAX_ASSISTANT_LENGTH), | |
| temperature=0.7, | |
| do_sample=True, | |
| pad_token_id=tokenizer.pad_token_id, | |
| bos_token_id=tokenizer.bos_token_id, | |
| eos_token_id=tokenizer.eos_token_id, | |
| bad_words_ids=[[tokenizer.unk_token_id]] | |
| ) | |
| output = tokenizer.decode(output_ids.tolist()[0]) | |
| formatted_output_all = format_output(output) | |
| response = f"Assistant:{formatted_output_all.split('応答:')[-1].strip()}" | |
| conversation_history.append(f"User:{new_conversation}".replace("\n", "\\n")) | |
| conversation_history.append(response.replace("\n", "\\n")) | |
| return formatted_output_all, response | |
| def home(): | |
| """ホームページをレンダリング""" | |
| return render_template_string(HTML_TEMPLATE) | |
| def generate(): | |
| """Flaskエンドポイント: /generate""" | |
| data = request.json | |
| role_instruction = data.get('role_instruction', []) | |
| conversation_history = data.get('conversation_history', []) | |
| new_conversation = data.get('new_conversation', "") | |
| if not role_instruction or not new_conversation: | |
| return jsonify({"error": "role_instruction and new_conversation are required fields"}), 400 | |
| formatted_output_all, response = generate_response(role_instruction, conversation_history, new_conversation) | |
| # ここでSSEを介してリアルタイムで応答をストリームします | |
| for word in response.split(): | |
| sse.publish({"message": word}, type='message') | |
| time.sleep(0.5) # 送信間隔をシミュレート | |
| return jsonify({"response": response, "conversation_history": conversation_history}) | |
| if __name__ == '__main__': | |
| app.run(debug=True, host="0.0.0.0", port=7860) | |