Yurii Paniv commited on
Commit
f0666ff
·
1 Parent(s): 6a2e815

Add zero gpu support

Browse files
Files changed (1) hide show
  1. app.py +2 -0
app.py CHANGED
@@ -5,6 +5,7 @@ from transformers import MistralForCausalLM, TextIteratorStreamer, AutoTokenizer
5
  from time import sleep
6
  from threading import Thread
7
  from torch import float16
 
8
 
9
  config = PeftConfig.from_pretrained("lang-uk/dragoman")
10
  quant_config = BitsAndBytesConfig(
@@ -20,6 +21,7 @@ model = MistralForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1",
20
  model = PeftModel.from_pretrained(model, "lang-uk/dragoman")
21
  tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1", use_fast=False, add_bos_token=False)
22
 
 
23
  def translate(input_text):
24
  # iteratively generate
25
  input_text = input_text.strip()
 
5
  from time import sleep
6
  from threading import Thread
7
  from torch import float16
8
+ import spaces
9
 
10
  config = PeftConfig.from_pretrained("lang-uk/dragoman")
11
  quant_config = BitsAndBytesConfig(
 
21
  model = PeftModel.from_pretrained(model, "lang-uk/dragoman")
22
  tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1", use_fast=False, add_bos_token=False)
23
 
24
+ @spaces.GPU(duration=200)
25
  def translate(input_text):
26
  # iteratively generate
27
  input_text = input_text.strip()