Update app.py

app.py

@@ -1,11 +1,29 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+import os
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
 """
-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
+Load model and tokenizer directly using transformers
 """
-client = InferenceClient("PlantWisdom/Data_Management_Mistral")
+model_name = "PlantWisdom/Data_Management_Mistral"
 
+# Configure quantization for lower memory usage
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.float16
+)
+
+# Load tokenizer and model
+print("Loading tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+print("Loading model...")
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    quantization_config=quantization_config,
+    device_map="auto",
+)
 
 def respond(
     message,
@@ -26,19 +44,29 @@ def respond(
     
     prompt += f"[INST] {message} [/INST]\n"
 
-    response = ""
+    # Print the prompt for debugging
+    print(f"Prompt: {prompt}")
 
-    for token in client.text_generation(
-        prompt,
+    # Encode the prompt
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    
+    # Generate tokens
+    print("Generating response...")
+    
+    # Generate without streaming for simplicity
+    generated_ids = model.generate(
+        inputs.input_ids,
         max_new_tokens=max_tokens,
-        stream=True,
+        do_sample=True,
         temperature=temperature,
         top_p=top_p,
-        do_sample=True,
-    ):
-        response += token
-        yield response
-
+        pad_token_id=tokenizer.eos_token_id,
+    )
+    
+    # Decode the response
+    full_response = tokenizer.decode(generated_ids[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
+    print(f"Response generated: {full_response[:50]}...")
+    return full_response
 
 """
 For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
@@ -46,19 +74,20 @@ For information on how to customize the ChatInterface, peruse the gradio docs: h
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[
-        gr.Textbox(value="You are a Microsoft 365 data management assistant.", 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.Textbox(value="You are a Microsoft 365 data management assistant specialized in SharePoint and OneDrive. Answer questions concisely and accurately.", label="System message"),
+        gr.Slider(minimum=1, maximum=1024, value=256, step=1, label="Max new tokens"),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature"),
         gr.Slider(
             minimum=0.1,
             maximum=1.0,
-            value=0.95,
+            value=0.9,
             step=0.05,
             label="Top-p (nucleus sampling)",
         ),
     ],
+    title="Microsoft 365 Data Management Assistant",
+    description="Ask questions about SharePoint, OneDrive, and other Microsoft 365 data management topics."
 )
 
-
 if __name__ == "__main__":
     demo.launch()