CreitinGameplays commited on
Commit
6250663
·
verified ·
1 Parent(s): 3d6b951

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -0
app.py CHANGED
@@ -1,5 +1,6 @@
1
  import gradio as gr
2
  import torch
 
3
  import bitsandbytes as bnb
4
  from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
5
 
@@ -20,6 +21,7 @@ model = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=bnb
20
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
21
  #model.to(device)
22
 
 
23
  def generate_text(user_prompt):
24
  """Generates text using the ConvAI model from Hugging Face Transformers and removes the user prompt."""
25
  # Construct the full prompt with system introduction, user prompt, and assistant role
 
1
  import gradio as gr
2
  import torch
3
+ import spaces
4
  import bitsandbytes as bnb
5
  from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
6
 
 
21
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
22
  #model.to(device)
23
 
24
+ @spaces.GPU(duration=120)
25
  def generate_text(user_prompt):
26
  """Generates text using the ConvAI model from Hugging Face Transformers and removes the user prompt."""
27
  # Construct the full prompt with system introduction, user prompt, and assistant role