kwabs22 commited on
Commit
ba8ad86
·
1 Parent(s): b66d017

Will it load?

Browse files
Files changed (1) hide show
  1. app.py +2 -0
app.py CHANGED
@@ -1,6 +1,7 @@
1
  import gradio as gr
2
  from transformers import AutoTokenizer, AutoModelForCausalLM
3
  import torch
 
4
 
5
  # Load model and tokenizer
6
  tokenizer = AutoTokenizer.from_pretrained("ISTA-DASLab/Meta-Llama-3.1-70B-AQLM-PV-2Bit-1x16")
@@ -8,6 +9,7 @@ model = AutoModelForCausalLM.from_pretrained("ISTA-DASLab/Meta-Llama-3.1-70B-AQL
8
  model = model.to('cuda') # Move the model to GPU if available
9
 
10
  # Define a function for generating text from a prompt
 
11
  def generate_text(prompt):
12
  inputs = tokenizer(prompt, return_tensors="pt").to('cuda') # Tokenize input and move to GPU
13
  outputs = model.generate(inputs.input_ids, max_length=100) # Generate output text
 
1
  import gradio as gr
2
  from transformers import AutoTokenizer, AutoModelForCausalLM
3
  import torch
4
+ import spaces
5
 
6
  # Load model and tokenizer
7
  tokenizer = AutoTokenizer.from_pretrained("ISTA-DASLab/Meta-Llama-3.1-70B-AQLM-PV-2Bit-1x16")
 
9
  model = model.to('cuda') # Move the model to GPU if available
10
 
11
  # Define a function for generating text from a prompt
12
+ @spaces.GPU
13
  def generate_text(prompt):
14
  inputs = tokenizer(prompt, return_tensors="pt").to('cuda') # Tokenize input and move to GPU
15
  outputs = model.generate(inputs.input_ids, max_length=100) # Generate output text