david-thrower commited on
Commit
b6fbedf
·
verified ·
1 Parent(s): c373b8a

Update app.py

Browse files

Add spaces.GPU decorator

Files changed (1) hide show
  1. app.py +8 -2
app.py CHANGED
@@ -2,6 +2,7 @@ from time import sleep
2
 
3
  import gradio as gr
4
  from transformers import AutoModelForCausalLM, AutoTokenizer
 
5
  import torch
6
  from duckduckgo_search import DDGS
7
  import re
@@ -322,6 +323,11 @@ class Applicant:
322
  def __init__(self, resume):
323
  self.resume = resume
324
 
 
 
 
 
 
325
  def smol_lm_jd_process(job_description, system_prompt, max_new_tokens=512):
326
  prompt = f"""<|im_start|>system
327
  {system_prompt}<|im_end|>
@@ -329,8 +335,8 @@ def smol_lm_jd_process(job_description, system_prompt, max_new_tokens=512):
329
  {job_description}<|im_end|>
330
  <|im_start|>assistant
331
  """
332
- inputs = tokenizer(prompt, return_tensors="pt").to(device)
333
- output = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=True, temperature=0.6, top_k=40, top_p=0.9, repetition_penalty=1.1)
334
  response = tokenizer.decode(output[0], skip_special_tokens=False)
335
  start_idx = response.find("<|im_start|>assistant")
336
  end_idx = response.find("<|im_end|>", start_idx)
 
2
 
3
  import gradio as gr
4
  from transformers import AutoModelForCausalLM, AutoTokenizer
5
+ import spaces
6
  import torch
7
  from duckduckgo_search import DDGS
8
  import re
 
323
  def __init__(self, resume):
324
  self.resume = resume
325
 
326
+ @spaces.GPU
327
+ def write(inputs):
328
+ _output = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=True, temperature=0.6, top_k=40, top_p=0.9, repetition_penalty=1.1)
329
+ return _output
330
+
331
  def smol_lm_jd_process(job_description, system_prompt, max_new_tokens=512):
332
  prompt = f"""<|im_start|>system
333
  {system_prompt}<|im_end|>
 
335
  {job_description}<|im_end|>
336
  <|im_start|>assistant
337
  """
338
+ inputs = tokenizer(prompt, return_tensors="pt").to(device)
339
+ output = write(inputs)
340
  response = tokenizer.decode(output[0], skip_special_tokens=False)
341
  start_idx = response.find("<|im_start|>assistant")
342
  end_idx = response.find("<|im_end|>", start_idx)