Update app.py

app.py

@@ -1,6 +1,18 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+from threading import Thread
+import torch
 
+model_id = "Molchevsky/ai_resume_llama-3.2-3b"
+
+# Load model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    torch_dtype=torch.float16,  # Use float16 to save memory
+    device_map="cpu",  # Explicitly set to CPU
+    low_cpu_mem_usage=True  # Helps with memory on load
+)
 
 def respond(
     message,
@@ -9,36 +21,35 @@ def respond(
     max_tokens,
     temperature,
     top_p,
-    hf_token: gr.OAuthToken,
 ):
-    """
-    For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-    """
-    client = InferenceClient(token=hf_token.token, model="Molchevsky/ai_resume_llama-3.2-3b")
-
     messages = [{"role": "system", "content": system_message}]
 
     messages.extend(history)
 
     messages.append({"role": "user", "content": message})
 
-    response = ""
+    # Apply chat template
+    input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(model.device)
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        choices = message.choices
-        token = ""
-        if len(choices) and choices[0].delta.content:
-            token = choices[0].delta.content
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
 
-        response += token
-        yield response
+    generation_kwargs = {
+        "input_ids": input_ids,
+        "streamer": streamer,
+        "max_new_tokens": max_tokens,
+        "temperature": temperature,
+        "top_p": top_p,
+        "do_sample": True if temperature > 0 else False,
+    }
 
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+
+    response = ""
+    for new_token in streamer:
+        response += new_token
+        yield response
 
 """
 For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
@@ -61,10 +72,7 @@ chatbot = gr.ChatInterface(
 )
 
 with gr.Blocks() as demo:
-    with gr.Sidebar():
-        gr.LoginButton()
     chatbot.render()
 
-
 if __name__ == "__main__":
-    demo.launch(debug=True)
+    demo.launch(debug=True)