hanzla commited on
Commit
d24a963
1 Parent(s): 09e1b8b
Files changed (1) hide show
  1. app.py +8 -4
app.py CHANGED
@@ -2,13 +2,17 @@ import gradio as gr
2
  import spaces
3
  import torch
4
 
 
 
 
 
5
 
 
 
 
6
  @spaces.GPU
7
  def yes_man(message, history):
8
- if message.endswith("?"):
9
- return "Yes"
10
- else:
11
- return "Ask me anything!"
12
 
13
  gr.ChatInterface(
14
  yes_man,
 
2
  import spaces
3
  import torch
4
 
5
+ import transformers
6
+ import torch
7
+
8
+ model_id = "meta-llama/Meta-Llama-3-8B"
9
 
10
+ pipeline = transformers.pipeline(
11
+ "text-generation", model=model_id, model_kwargs={"torch_dtype": torch.bfloat16}, device_map="auto"
12
+ )
13
  @spaces.GPU
14
  def yes_man(message, history):
15
+ return pipeline(message)
 
 
 
16
 
17
  gr.ChatInterface(
18
  yes_man,