OzoneAsai commited on
Commit
22dd958
1 Parent(s): f2df7fd

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +126 -35
app.py CHANGED
@@ -1,57 +1,141 @@
1
- from flask import Flask, request, jsonify
2
- from transformers import AutoTokenizer, AutoModelForCausalLM
3
  import torch
 
 
 
 
 
4
 
 
5
  app = Flask(__name__)
 
 
6
 
7
- # モデルとトークナイザーの読み込み
8
- tokenizer = AutoTokenizer.from_pretrained("inu-ai/alpaca-guanaco-japanese-gpt-1b", use_fast=False)
9
- model = AutoModelForCausalLM.from_pretrained("inu-ai/alpaca-guanaco-japanese-gpt-1b")
10
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
11
- model.to(device)
12
 
 
 
 
 
 
13
  MAX_ASSISTANT_LENGTH = 100
14
  MAX_INPUT_LENGTH = 1024
15
  INPUT_PROMPT = r'<s>\n以下は、タスクを説明する指示と、文脈のある入力の組み合わせです。要求を適切に満たす応答を書きなさい。\n[SEP]\n指示:\n{instruction}\n[SEP]\n入力:\n{input}\n[SEP]\n応答:\n'
16
  NO_INPUT_PROMPT = r'<s>\n以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。\n[SEP]\n指示:\n{instruction}\n[SEP]\n応答:\n'
17
 
18
- def get_default_role_instruction():
19
- return [
20
- "User:睡眠に悩む高校生です",
21
- "Assistant:では、お手伝いしましょう。!"
22
- ]
23
-
24
- def get_default_conversation_history():
25
- return [
26
- "User: こんにちは、今日は一日を有効に使いたいのですが、何かアドバイスはありますか?",
27
- "Assistant: こんにちは!一日の計画を立てることはとても重要です。朝、昼、晩それぞれの時間帯でやるべきことをリストにまとめると良いですよ。",
28
- "User: なるほど、具体的にはどんな内容をリストに入れればいいですか?",
29
- "Assistant: 朝は、起床後のルーチンや朝食の準備、健康的な運動を入れると良いですね。昼は、仕事や勉強の計画、休憩時間、昼食の準備が考えられます。晩は、夕食の準備や家事、リラックスタイム、就寝前のルーチンなどが含まれます。"
30
- ]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
31
 
32
  def prepare_input(role_instruction, conversation_history, new_conversation):
 
33
  instruction = "".join([f"{text}\n" for text in role_instruction])
34
  instruction += "\n".join(conversation_history)
35
  input_text = f"User:{new_conversation}"
36
  return INPUT_PROMPT.format(instruction=instruction, input=input_text)
37
 
38
  def format_output(output):
 
39
  return output.lstrip("<s>").rstrip("</s>").replace("[SEP]", "").replace("\\n", "\n")
40
 
41
  def trim_conversation_history(conversation_history, max_length):
 
42
  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:
43
  conversation_history.pop(0)
44
  conversation_history.pop(0)
45
  return conversation_history
46
 
47
  def generate_response(role_instruction, conversation_history, new_conversation):
 
48
  conversation_history = trim_conversation_history(conversation_history, MAX_ASSISTANT_LENGTH)
49
  input_text = prepare_input(role_instruction, conversation_history, new_conversation)
50
  token_ids = tokenizer.encode(input_text, add_special_tokens=False, return_tensors="pt")
51
 
52
  with torch.no_grad():
53
  output_ids = model.generate(
54
- token_ids.to(device),
55
  min_length=len(token_ids[0]),
56
  max_length=min(MAX_INPUT_LENGTH, len(token_ids[0]) + MAX_ASSISTANT_LENGTH),
57
  temperature=0.7,
@@ -71,23 +155,30 @@ def generate_response(role_instruction, conversation_history, new_conversation):
71
 
72
  return formatted_output_all, response
73
 
74
- @app.route('/api', methods=['POST'])
75
- def api():
76
- try:
77
- data = request.json
78
- role_instruction = data.get('role_instruction', get_default_role_instruction())
79
- conversation_history = data.get('conversation_history', get_default_conversation_history())
80
- new_conversation = data.get('new_conversation', "")
 
 
 
 
 
81
 
82
- if not new_conversation:
83
- return jsonify({"error": "new_conversation is required"}), 400
84
 
85
- formatted_output_all, response = generate_response(role_instruction, conversation_history, new_conversation)
86
- return jsonify({"response": response, "conversation_history": conversation_history})
 
 
 
 
87
 
88
- except Exception as e:
89
- print(f"Error: {e}")
90
- return jsonify({"error": "An error occurred"}), 500
91
 
92
  if __name__ == '__main__':
93
- app.run(host='0.0.0.0', port=7860)
 
 
 
1
  import torch
2
+ from transformers import AutoTokenizer, AutoModelForCausalLM
3
+ from flask import Flask, request, jsonify, render_template_string
4
+ import time
5
+ from flask_sse import sse
6
+ import redis
7
 
8
+ # Flaskアプリケーションの設定
9
  app = Flask(__name__)
10
+ app.config["REDIS_URL"] = "redis://localhost:6379/0"
11
+ app.register_blueprint(sse, url_prefix='/stream')
12
 
13
+ # デバイスの設定
 
 
14
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
15
 
16
+ # トークナイザーとモデルの読み込み
17
+ tokenizer = AutoTokenizer.from_pretrained("inu-ai/alpaca-guanaco-japanese-gpt-1b", use_fast=False)
18
+ model = AutoModelForCausalLM.from_pretrained("inu-ai/alpaca-guanaco-japanese-gpt-1b").to(device)
19
+
20
+ # 定数
21
  MAX_ASSISTANT_LENGTH = 100
22
  MAX_INPUT_LENGTH = 1024
23
  INPUT_PROMPT = r'<s>\n以下は、タスクを説明する指示と、文脈のある入力の組み合わせです。要求を適切に満たす応答を書きなさい。\n[SEP]\n指示:\n{instruction}\n[SEP]\n入力:\n{input}\n[SEP]\n応答:\n'
24
  NO_INPUT_PROMPT = r'<s>\n以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。\n[SEP]\n指示:\n{instruction}\n[SEP]\n応答:\n'
25
 
26
+ # HTMLテンプレート
27
+ HTML_TEMPLATE = """
28
+ <!DOCTYPE html>
29
+ <html lang="ja">
30
+ <head>
31
+ <meta charset="UTF-8">
32
+ <meta name="viewport" content="width=device-width, initial-scale=1.0">
33
+ <title>Chat Interface</title>
34
+ <style>
35
+ body { font-family: Arial, sans-serif; }
36
+ .container { max-width: 600px; margin: auto; padding: 20px; }
37
+ .chat-box { border: 1px solid #ccc; padding: 10px; height: 300px; overflow-y: scroll; }
38
+ .chat-entry { margin-bottom: 10px; }
39
+ .chat-entry.user { text-align: right; }
40
+ .input-group { display: flex; }
41
+ .input-group input { flex: 1; padding: 10px; border: 1px solid #ccc; border-radius: 4px; }
42
+ .input-group button { padding: 10px; border: none; background-color: #28a745; color: white; cursor: pointer; }
43
+ .input-group button:hover { background-color: #218838; }
44
+ </style>
45
+ </head>
46
+ <body>
47
+ <div class="container">
48
+ <h1>Chat Interface</h1>
49
+ <div class="chat-box" id="chat-box"></div>
50
+ <div class="input-group">
51
+ <input type="text" id="user-input" placeholder="質問を入力してください...">
52
+ <button onclick="sendMessage()">送信</button>
53
+ </div>
54
+ </div>
55
+ <script>
56
+ const chatBox = document.getElementById('chat-box');
57
+ const userInput = document.getElementById('user-input');
58
+ let conversationHistory = [];
59
+
60
+ function addMessageToChat(role, message) {
61
+ const entry = document.createElement('div');
62
+ entry.className = 'chat-entry ' + role;
63
+ entry.textContent = role === 'user' ? 'User: ' + message : 'Assistant: ' + message;
64
+ chatBox.appendChild(entry);
65
+ chatBox.scrollTop = chatBox.scrollHeight;
66
+ }
67
+
68
+ function sendMessage() {
69
+ const message = userInput.value.trim();
70
+ if (message === '') return;
71
+
72
+ addMessageToChat('user', message);
73
+ userInput.value = '';
74
+
75
+ fetch('/generate', {
76
+ method: 'POST',
77
+ headers: { 'Content-Type': 'application/json' },
78
+ body: JSON.stringify({
79
+ role_instruction: [
80
+ "User:あなたは「ずんだもん」なのだ。東北ずん子の武器である「ずんだアロー」に変身する妖精またはマスコットなのだ。一人称は「ボク」で語尾に「なのだ」を付けてしゃべるのだ���",
81
+ "Assistant:了解したのだ!"
82
+ ],
83
+ conversation_history: conversationHistory,
84
+ new_conversation: message
85
+ })
86
+ })
87
+ .then(response => response.json())
88
+ .then(data => {
89
+ const assistantMessage = data.response.split('Assistant:')[1].trim();
90
+ addMessageToChat('assistant', assistantMessage);
91
+ conversationHistory.push('User:' + message);
92
+ conversationHistory.push('Assistant:' + assistantMessage);
93
+ })
94
+ .catch(error => {
95
+ console.error('Error:', error);
96
+ alert('エラーが発生しました。コンソールを確認してください。');
97
+ });
98
+ }
99
+
100
+ // SSEの設定
101
+ const eventSource = new EventSource("/stream");
102
+
103
+ eventSource.onmessage = function(event) {
104
+ const message = event.data;
105
+ addMessageToChat('assistant', message);
106
+ };
107
+ </script>
108
+ </body>
109
+ </html>
110
+ """
111
 
112
  def prepare_input(role_instruction, conversation_history, new_conversation):
113
+ """入力テキストを整形する関数"""
114
  instruction = "".join([f"{text}\n" for text in role_instruction])
115
  instruction += "\n".join(conversation_history)
116
  input_text = f"User:{new_conversation}"
117
  return INPUT_PROMPT.format(instruction=instruction, input=input_text)
118
 
119
  def format_output(output):
120
+ """生成された出力を整形する関数"""
121
  return output.lstrip("<s>").rstrip("</s>").replace("[SEP]", "").replace("\\n", "\n")
122
 
123
  def trim_conversation_history(conversation_history, max_length):
124
+ """会話履歴を最大長に収めるために調整する関数"""
125
  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:
126
  conversation_history.pop(0)
127
  conversation_history.pop(0)
128
  return conversation_history
129
 
130
  def generate_response(role_instruction, conversation_history, new_conversation):
131
+ """新しい会話に対する応答を生成する関数"""
132
  conversation_history = trim_conversation_history(conversation_history, MAX_ASSISTANT_LENGTH)
133
  input_text = prepare_input(role_instruction, conversation_history, new_conversation)
134
  token_ids = tokenizer.encode(input_text, add_special_tokens=False, return_tensors="pt")
135
 
136
  with torch.no_grad():
137
  output_ids = model.generate(
138
+ token_ids.to(model.device),
139
  min_length=len(token_ids[0]),
140
  max_length=min(MAX_INPUT_LENGTH, len(token_ids[0]) + MAX_ASSISTANT_LENGTH),
141
  temperature=0.7,
 
155
 
156
  return formatted_output_all, response
157
 
158
+ @app.route('/')
159
+ def home():
160
+ """ホームページをレンダリング"""
161
+ return render_template_string(HTML_TEMPLATE)
162
+
163
+ @app.route('/generate', methods=['POST'])
164
+ def generate():
165
+ """Flaskエンドポイント: /generate"""
166
+ data = request.json
167
+ role_instruction = data.get('role_instruction', [])
168
+ conversation_history = data.get('conversation_history', [])
169
+ new_conversation = data.get('new_conversation', "")
170
 
171
+ if not role_instruction or not new_conversation:
172
+ return jsonify({"error": "role_instruction and new_conversation are required fields"}), 400
173
 
174
+ formatted_output_all, response = generate_response(role_instruction, conversation_history, new_conversation)
175
+
176
+ # ここでSSEを介してリアルタイムで応答をストリームします
177
+ for word in response.split():
178
+ sse.publish({"message": word}, type='message')
179
+ time.sleep(0.5) # 送信間隔をシミュレート
180
 
181
+ return jsonify({"response": response, "conversation_history": conversation_history})
 
 
182
 
183
  if __name__ == '__main__':
184
+ app.run(debug=True, host="0.0.0.0", port=7860)