ndwdgda commited on
Commit
b77a8b6
1 Parent(s): 028ec90

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +12 -18
app.py CHANGED
@@ -1,11 +1,8 @@
1
  import gradio as gr
2
- from huggingface_hub import InferenceClient
3
-
4
- """
5
- For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
6
- """
7
- client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
8
 
 
 
9
 
10
  def respond(
11
  message,
@@ -27,21 +24,19 @@ def respond(
27
 
28
  response = ""
29
 
30
- for message in client.chat_completion(
31
- messages,
32
- max_tokens=max_tokens,
33
- stream=True,
 
34
  temperature=temperature,
35
  top_p=top_p,
36
- ):
37
- token = message.choices[0].delta.content
38
 
39
- response += token
40
- yield response
41
 
42
- """
43
- For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
44
- """
45
  demo = gr.ChatInterface(
46
  respond,
47
  additional_inputs=[
@@ -58,6 +53,5 @@ demo = gr.ChatInterface(
58
  ],
59
  )
60
 
61
-
62
  if __name__ == "__main__":
63
  demo.launch()
 
1
  import gradio as gr
2
+ from transformers import pipeline
 
 
 
 
 
3
 
4
+ # Initialize the text generation pipeline with the specified model
5
+ pipe = pipeline("text-generation", model="chargoddard/Yi-34B-Llama", device=0)
6
 
7
  def respond(
8
  message,
 
24
 
25
  response = ""
26
 
27
+ # Generate the response using the pipeline
28
+ result = pipe(
29
+ messages[-1]["content"],
30
+ max_length=max_tokens,
31
+ num_return_sequences=1,
32
  temperature=temperature,
33
  top_p=top_p,
34
+ )
 
35
 
36
+ response = result[0]['generated_text']
37
+ yield response
38
 
39
+ # Gradio interface setup
 
 
40
  demo = gr.ChatInterface(
41
  respond,
42
  additional_inputs=[
 
53
  ],
54
  )
55
 
 
56
  if __name__ == "__main__":
57
  demo.launch()