cody82 commited on
Commit
000988a
·
verified ·
1 Parent(s): 4436ef3

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +16 -23
app.py CHANGED
@@ -1,13 +1,13 @@
1
  import os
2
  os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
3
 
 
 
4
  import torch
5
  from transformers import AutoTokenizer, AutoModelForCausalLM
6
- import gradio as gr
7
- from fastapi import FastAPI
8
  import uvicorn
9
 
10
- # === Модель ===
11
  model_id = "sberbank-ai/rugpt3medium_based_on_gpt2"
12
  tokenizer = AutoTokenizer.from_pretrained(model_id)
13
  model = AutoModelForCausalLM.from_pretrained(model_id)
@@ -15,14 +15,23 @@ model = AutoModelForCausalLM.from_pretrained(model_id)
15
  device = "cuda" if torch.cuda.is_available() else "cpu"
16
  model.to(device)
17
 
 
18
  context = (
19
  "Университет Иннополис был основан в 2012 году. "
20
  "Это современный вуз в России, специализирующийся на IT и робототехнике, "
21
  "расположенный в городе Иннополис, Татарстан.\n"
22
  )
23
 
24
- def respond(message, history=None):
25
- prompt = f"Прочитай текст и ответь на вопрос:\n\n{context}\n\nВопрос: {message}\nОтвет:"
 
 
 
 
 
 
 
 
26
  input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
27
 
28
  with torch.no_grad():
@@ -42,24 +51,8 @@ def respond(message, history=None):
42
  else:
43
  answer = output[len(prompt):].strip()
44
 
45
- return answer
46
-
47
- # === Gradio интерфейс (обязательно должен быть `demo` для Hugging Face Spaces) ===
48
- demo = gr.ChatInterface(fn=respond, title="Иннополис Бот")
49
-
50
- # === FastAPI (опционально, если нужен API) ===
51
- app = FastAPI()
52
-
53
- @app.get("/health")
54
- def health_check():
55
- return {"status": "OK"}
56
-
57
- @app.post("/ask")
58
- async def ask(question: str):
59
- return {"answer": respond(question)}
60
 
61
- # === Если запускаем локально (не в Spaces) ===
62
  if __name__ == "__main__":
63
- # Для локального теста с API
64
- app = gr.mount_gradio_app(app, demo, path="/")
65
  uvicorn.run(app, host="0.0.0.0", port=8000)
 
1
  import os
2
  os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
3
 
4
+ from fastapi import FastAPI
5
+ from pydantic import BaseModel
6
  import torch
7
  from transformers import AutoTokenizer, AutoModelForCausalLM
 
 
8
  import uvicorn
9
 
10
+ # === Загрузка модели ===
11
  model_id = "sberbank-ai/rugpt3medium_based_on_gpt2"
12
  tokenizer = AutoTokenizer.from_pretrained(model_id)
13
  model = AutoModelForCausalLM.from_pretrained(model_id)
 
15
  device = "cuda" if torch.cuda.is_available() else "cpu"
16
  model.to(device)
17
 
18
+ # Контекст для модели
19
  context = (
20
  "Университет Иннополис был основан в 2012 году. "
21
  "Это современный вуз в России, специализирующийся на IT и робототехнике, "
22
  "расположенный в городе Иннополис, Татарстан.\n"
23
  )
24
 
25
+ # === FastAPI приложение ===
26
+ app = FastAPI()
27
+
28
+ class QuestionRequest(BaseModel):
29
+ question: str
30
+
31
+ @app.post("/ask")
32
+ def generate_answer(request: QuestionRequest):
33
+ """Обрабатывает POST-запрос с вопросом и возвращает ответ модели."""
34
+ prompt = f"Прочитай текст и ответь на вопрос:\n\n{context}\n\nВопрос: {request.question}\nОтвет:"
35
  input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
36
 
37
  with torch.no_grad():
 
51
  else:
52
  answer = output[len(prompt):].strip()
53
 
54
+ return {"answer": answer}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
55
 
56
+ # Точка входа для запуска сервера
57
  if __name__ == "__main__":
 
 
58
  uvicorn.run(app, host="0.0.0.0", port=8000)