Update app.py

Added inference endpoint.

app.py

@@ -1,8 +1,10 @@
 from time import sleep
+from os import getenv
 
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import spaces
+from openai import OpenAI
 import torch
 from duckduckgo_search import DDGS
 import re
@@ -15,6 +17,8 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 
+
+
 class DDGSSearchClient:
     def __init__(self, max_retries=4, timeout=35, backoff_factor=1):
         """
@@ -323,10 +327,40 @@ class Applicant:
     def __init__(self, resume):
         self.resume = resume
 
-@spaces.GPU
-def write(inputs, max_new_tokens, do_sample=True, temperature=0.6, top_k=40, top_p=0.9, repetition_penalty=1.1):
-    _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)
-    return _output
+# @spaces.GPU
+# def write(inputs, max_new_tokens, do_sample=True, temperature=0.6, top_k=40, top_p=0.9, repetition_penalty=1.1):
+#     _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)
+#     return _output
+
+def writing_task(prompt: str) -> str:
+    api_key = getenv("HF_TOKEN")
+    if not api_key:
+        raise ValueError("Huggingface token missing. Need to set HF_TOKEN, refer to https://discuss.huggingface.co/t/how-to-manage-user-secrets-and-api-keys/67948")
+    client = OpenAI(
+        base_url="https://router.huggingface.co/v1",
+        api_key = getenv("HF_TOKEN")
+    )
+
+    completion = client.chat.completions.create(
+        model="HuggingFaceTB/SmolLM3-3B:hf-inference",
+        messages=[
+            {
+                "role": "user",
+                "content": prompt
+            }
+        ],
+    )
+
+    raw_response_content = completion.choices[0].message.content
+    content_split = raw_response_content.split("</think>")
+    if len(content_split) > 1:
+        think = content_split[0]
+        content = "".join(content_split[1:])
+    else:
+        think = content_split[0]
+        content = "No data found."
+    
+    return content
 
 def smol_lm_jd_process(job_description, system_prompt, max_new_tokens=512):
     prompt = f"""<|im_start|>system
@@ -335,12 +369,13 @@ def smol_lm_jd_process(job_description, system_prompt, max_new_tokens=512):
 {job_description}<|im_end|>
 <|im_start|>assistant
 """
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)   
-    output = write(inputs, max_new_tokens=max_new_tokens)
-    response = tokenizer.decode(output[0], skip_special_tokens=False)
-    start_idx = response.find("<|im_start|>assistant")
-    end_idx = response.find("<|im_end|>", start_idx)
-    response = response[start_idx + len("<|im_start|>assistant\n"):end_idx].strip()
+    # inputs = tokenizer(prompt, return_tensors="pt").to(device)   
+    # output = write(inputs, max_new_tokens=max_new_tokens)
+    # response = tokenizer.decode(output[0], skip_special_tokens=False)
+    # start_idx = response.find("<|im_start|>assistant")
+    # end_idx = response.find("<|im_end|>", start_idx)
+    # response = response[start_idx + len("<|im_start|>assistant\n"):end_idx].strip()
+    response = writing_task(prompt)
     return response
 
 def process_job_description(company_name, company_url, job_description, resume):