Update

app.py

@@ -1,52 +1,64 @@
-import os
 import gradio as gr
-import torch
-from transformers import AutoTokenizer, AutoModel
-from safetensors.torch import load_file
+from huggingface_hub import InferenceClient
 
-# Load the Hugging Face API token from environment variable
-token = os.getenv("HUGGINGFACE_API_TOKEN")
-if not token:
-    raise ValueError("HUGGINGFACE_API_TOKEN is not set. Please add it in the Secrets section of your Space.")
+"""
+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("Grandediw/lora_model")
 
-# Configure device
-device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load the tokenizer and model using the token
-model_repo = "Grandediw/lora_model"
-tokenizer = AutoTokenizer.from_pretrained(model_repo, token=token)
-base_model = AutoModel.from_pretrained(model_repo, token=token)
+def respond(
+    message,
+    history: list[tuple[str, str]],
+    system_message,
+    max_tokens,
+    temperature,
+    top_p,
+):
+    messages = [{"role": "system", "content": system_message}]
 
-# Load the LoRA adapter weights
-lora_weights_path = "adapter_model.safetensors"  
-lora_weights = load_file(lora_weights_path)
+    for val in history:
+        if val[0]:
+            messages.append({"role": "user", "content": val[0]})
+        if val[1]:
+            messages.append({"role": "assistant", "content": val[1]})
 
-# Apply LoRA weights to the base model
-for name, param in base_model.named_parameters():
-    if name in lora_weights:
-        param.data += lora_weights[name].to(device, dtype=param.dtype)
+    messages.append({"role": "user", "content": message})
 
-# Move the model to the device
-base_model = base_model.to(device)
+    response = ""
 
-# Define the inference function
-def infer(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
-    outputs = base_model(**inputs)
-    # Placeholder return, modify based on your specific model task
-    return outputs.last_hidden_state.mean(dim=1).cpu().detach().numpy()
+    for message in client.chat_completion(
+        messages,
+        max_tokens=max_tokens,
+        stream=True,
+        temperature=temperature,
+        top_p=top_p,
+    ):
+        token = message.choices[0].delta.content
 
-# Gradio interface Update
-with gr.Blocks() as demo:
-    gr.Markdown("## LoRA Model Inference")
+        response += token
+        yield response
 
-    with gr.Row():
-        prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here...")
-        generate_button = gr.Button("Generate")
 
-    output = gr.Textbox(label="Output")
+"""
+For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
+"""
+demo = gr.ChatInterface(
+    respond,
+    additional_inputs=[
+        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
+        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
+        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
+        gr.Slider(
+            minimum=0.1,
+            maximum=1.0,
+            value=0.95,
+            step=0.05,
+            label="Top-p (nucleus sampling)",
+        ),
+    ],
+)
 
-    generate_button.click(fn=infer, inputs=[prompt], outputs=[output])
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()