cody82 commited on
Commit
5a42269
·
verified ·
1 Parent(s): c833cd6

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +10 -14
app.py CHANGED
@@ -1,24 +1,20 @@
1
- import gradio as gr
2
  import torch
 
3
  from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
4
- import spaces # обязательно для ZeroGPU
5
 
6
- model_name = "google/flan-t5-base"
7
- tokenizer = AutoTokenizer.from_pretrained(model_name)
8
- model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
9
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
10
- model.to(device)
11
 
12
  @spaces.GPU
13
  def respond(message, history=[]):
14
  prompt = f"Answer the question: {message}"
15
- inputs = tokenizer(prompt, return_tensors="pt").to(device)
16
- outputs = model.generate(**inputs, max_new_tokens=100)
17
- response = tokenizer.decode(outputs[0], skip_special_tokens=True)
18
  history.append((message, response))
19
  return history
20
 
21
- demo = gr.ChatInterface(fn=respond)
22
-
23
- if __name__ == "__main__":
24
- demo.launch()
 
 
1
  import torch
2
+ import gradio as gr
3
  from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
4
+ import spaces
5
 
6
+ model_id = "google/flan-t5-base"
7
+ tokenizer = AutoTokenizer.from_pretrained(model_id)
8
+ model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
9
+ model.to("cuda" if torch.cuda.is_available() else "cpu")
 
10
 
11
  @spaces.GPU
12
  def respond(message, history=[]):
13
  prompt = f"Answer the question: {message}"
14
+ inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
15
+ output = model.generate(**inputs, max_new_tokens=100)
16
+ response = tokenizer.decode(output[0], skip_special_tokens=True)
17
  history.append((message, response))
18
  return history
19
 
20
+ gr.ChatInterface(fn=respond).launch()