KoDer123 commited on
Commit
5ee1df2
·
verified ·
1 Parent(s): 01ac9e3

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +66 -25
app.py CHANGED
@@ -9,6 +9,36 @@ import gc
9
  import signal
10
  from contextlib import contextmanager
11
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
12
  # Настройка логирования
13
  logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
14
  logger = logging.getLogger(__name__)
@@ -239,31 +269,38 @@ def respond(
239
  logger.info(f"Генерируем ответ на запрос: '{message[:50]}...' (длина промпта: {len(full_prompt)})")
240
 
241
  try:
242
- # Используем таймаут для генерации
243
- with time_limit(generation_timeout):
244
- # Токенизация входных данных
245
- inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
246
-
247
- # Генерация ответа
248
- outputs = model.generate(
249
- inputs.input_ids,
250
- max_new_tokens=max_tokens,
251
- temperature=temperature,
252
- top_p=top_p,
253
- do_sample=True,
254
- pad_token_id=tokenizer.pad_token_id,
255
- )
256
-
257
- # Декодирование полного вывода
258
- generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
259
-
260
- # Извлекаем только часть после "Ассистент: "
261
- response_start = generated_text.rfind("Ассистент: ") + len("Ассистент: ")
262
- if response_start >= len("Ассистент: "): # Проверяем, что "Ассистент: " найден
263
- response = generated_text[response_start:].strip()
264
- else:
265
- # Если не найдено, возвращаем весь текст
266
- response = generated_text.strip()
 
 
 
 
 
 
 
267
 
268
  end_time = time.time()
269
  generation_time = end_time - start_time
@@ -277,6 +314,10 @@ def respond(
277
  except Exception as e:
278
  logger.error(f"Ошибка при генерации ответа: {str(e)}")
279
  return f"Произошла ошибка при генерации ответа: {str(e)}"
 
 
 
 
280
 
281
  # Настройка интерфейса Gradio
282
  with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
9
  import signal
10
  from contextlib import contextmanager
11
 
12
+ import threading
13
+ import time
14
+
15
+ # Класс для обработки таймаута без использования signal
16
+ class TimeoutManager:
17
+ def __init__(self, seconds):
18
+ self.seconds = seconds
19
+ self.timeout_occurred = False
20
+ self.timer = None
21
+
22
+ def start(self):
23
+ self.timeout_occurred = False
24
+ self.timer = threading.Timer(self.seconds, self._timeout)
25
+ self.timer.daemon = True
26
+ self.timer.start()
27
+
28
+ def _timeout(self):
29
+ self.timeout_occurred = True
30
+
31
+ def stop(self):
32
+ if self.timer:
33
+ self.timer.cancel()
34
+
35
+ def check_timeout(self):
36
+ if self.timeout_occurred:
37
+ raise TimeoutException("Timeout occurred")
38
+
39
+ class TimeoutException(Exception):
40
+ pass
41
+
42
  # Настройка логирования
43
  logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
44
  logger = logging.getLogger(__name__)
 
269
  logger.info(f"Генерируем ответ на запрос: '{message[:50]}...' (длина промпта: {len(full_prompt)})")
270
 
271
  try:
272
+ # Настраиваем таймаут
273
+ timeout_mgr = TimeoutManager(generation_timeout)
274
+ timeout_mgr.start()
275
+
276
+ # Токенизация входных данных
277
+ inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
278
+
279
+ # Генерация ответа с периодической проверкой таймаута
280
+ gen_kwargs = {
281
+ "input_ids": inputs.input_ids,
282
+ "max_new_tokens": max_tokens,
283
+ "temperature": temperature,
284
+ "top_p": top_p,
285
+ "do_sample": True,
286
+ "pad_token_id": tokenizer.pad_token_id,
287
+ }
288
+
289
+ outputs = model.generate(**gen_kwargs)
290
+
291
+ # Останавливаем таймаут
292
+ timeout_mgr.stop()
293
+
294
+ # Декодирование полного вывода
295
+ generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
296
+
297
+ # Извлекаем только часть после "Ассистент: "
298
+ response_start = generated_text.rfind("Ассистент: ") + len("Ассистент: ")
299
+ if response_start >= len("Ассистент: "): # Проверяем, что "Ассистент: " найден
300
+ response = generated_text[response_start:].strip()
301
+ else:
302
+ # Если не найдено, возвращаем весь текст
303
+ response = generated_text.strip()
304
 
305
  end_time = time.time()
306
  generation_time = end_time - start_time
 
314
  except Exception as e:
315
  logger.error(f"Ошибка при генерации ответа: {str(e)}")
316
  return f"Произошла ошибка при генерации ответа: {str(e)}"
317
+ finally:
318
+ # Гарантируем остановку таймера
319
+ if 'timeout_mgr' in locals():
320
+ timeout_mgr.stop()
321
 
322
  # Настройка интерфейса Gradio
323
  with gr.Blocks(theme=gr.themes.Soft()) as demo: