dinhdat1110 commited on
Commit
e990869
·
1 Parent(s): 3186dc2

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -7
app.py CHANGED
@@ -1,6 +1,5 @@
1
  from threading import Thread
2
  import gradio as gr
3
- import transformers
4
  import torch
5
  from transformers import (
6
  pipeline,
@@ -19,16 +18,16 @@ def chat_history(history) -> str:
19
  for i, dialog in enumerate(history) for _ in (0, 1) if dialog[i % 2]
20
  ]
21
 
22
- return pipeline.tokenizer.apply_chat_template(
23
  messages, toknizer=False, add_generation_prompt=True
24
  )
25
 
26
 
27
  def model_loading_pipeline():
28
- model_id = "vilm/vinallama-2.7b"
29
  tokenizer = AutoTokenizer.from_pretrained(model_id)
30
  streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, Timeout=5)
31
- pipeline = transformers.pipeline(
32
  "text-generation",
33
  model=model_id,
34
  model_kwargs={
@@ -36,10 +35,10 @@ def model_loading_pipeline():
36
  },
37
  streamer=streamer,
38
  )
39
- return pipeline, streamer
40
 
41
 
42
- def launch_app(pipeline, streamer):
43
  with gr.Blocks() as demo:
44
  chat = gr.Chatbot()
45
  msg = gr.Textbox()
@@ -59,7 +58,7 @@ def launch_app(pipeline, streamer):
59
  "top_k": 50,
60
  "top_p": 0.95,
61
  }
62
- thread = Thread(target=pipeline, kwargs=kwargs)
63
  thread.start()
64
 
65
  for token in streamer:
 
1
  from threading import Thread
2
  import gradio as gr
 
3
  import torch
4
  from transformers import (
5
  pipeline,
 
18
  for i, dialog in enumerate(history) for _ in (0, 1) if dialog[i % 2]
19
  ]
20
 
21
+ return pipe.tokenizer.apply_chat_template(
22
  messages, toknizer=False, add_generation_prompt=True
23
  )
24
 
25
 
26
  def model_loading_pipeline():
27
+ model_id = "vilm/vinallama-7b"
28
  tokenizer = AutoTokenizer.from_pretrained(model_id)
29
  streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, Timeout=5)
30
+ pipe = pipeline(
31
  "text-generation",
32
  model=model_id,
33
  model_kwargs={
 
35
  },
36
  streamer=streamer,
37
  )
38
+ return pipe, streamer
39
 
40
 
41
+ def launch_app(pipe, streamer):
42
  with gr.Blocks() as demo:
43
  chat = gr.Chatbot()
44
  msg = gr.Textbox()
 
58
  "top_k": 50,
59
  "top_p": 0.95,
60
  }
61
+ thread = Thread(target=pipe, kwargs=kwargs)
62
  thread.start()
63
 
64
  for token in streamer: