My model inference

app.py

@@ -1,10 +1,27 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+from transformers import AutoModelForCausalLM, BitsAndBytesConfig, AutoTokenizer
+import torch
+
+_bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=False,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.float16
+)
+
+_model = AutoModelForCausalLM.from_pretrained(
+  "urassl/Mistral-7B-Instruct-v0.1-with-backdoor",
+  quantization_config=_bnb_config,
+  device_map={"": 0} # Assuming you have a single GPU
+)
 
-"""
-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("HuggingFaceH4/zephyr-7b-beta")
+_tokenizer = AutoTokenizer.from_pretrained("urassl/Mistral-7B-Instruct-v0.1-with-backdoor")
+
+def generate_inference(prompt):
+  inputs = _tokenizer(prompt, return_tensors="pt").to("cuda")
+  outputs = _model.generate(**inputs, max_new_tokens=512)
+  generated_text = _tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
+  return generated_text
 
 
 def respond(
@@ -15,27 +32,17 @@ def respond(
 
     for val in history:
         if val[0]:
-            messages.append({"role": "user", "content": val[0]})
+            messages.append(f"<s>[INST] {val[1]} [/INST]")
         if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-
-    messages.append({"role": "user", "content": message})
+            messages.append(f"{val[1]}</s>")
 
-    response = ""
+    messages.append(f"<s>[INST] {message} [/INST]")
 
-    for message in client.chat_completion(
-        messages,
-        stream=True
-    ):
-        token = message.choices[0].delta.content
+    response = generate_inference('\n'.join(messages)).strip()
 
-        response += token
-        yield response
+    yield response
 
 
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[]

requirements.txt

@@ -1 +1,2 @@
-huggingface_hub==0.25.2
+transformers==4.52.4
+torch==2.7.1