Update app.py

app.py

@@ -1,45 +1,58 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
+from peft import PeftModel
+import torch
 
 # -----------------------------
-# Respond function using text-generation
+# Load 4-bit Qwen model locally
 # -----------------------------
-async def respond(
-    message,
-    history: list[dict[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-    hf_token: gr.OAuthToken,
-):
-    """
-    Non-chat model version using text-generation.
-    Works for LoRA / 4-bit Qwen models.
-    """
-    client = InferenceClient(
-        token=hf_token.token,
-        model="unsloth/qwen2.5-math-1.5b-bnb-4bit"
-    )
+MODEL_NAME = "unsloth/qwen2.5-math-1.5b-bnb-4bit"
+
+print("Loading tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=False)
+
+print("Loading model in 4-bit...")
+base_model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    device_map="auto",
+    torch_dtype=torch.float16,
+    low_cpu_mem_usage=True
+)
+
+# Check if LoRA adapter exists
+try:
+    model = PeftModel.from_pretrained(base_model, MODEL_NAME, device_map="auto")
+except:
+    model = base_model
 
-    # Build prompt manually
+model.eval()
+
+# -----------------------------
+# Respond function
+# -----------------------------
+def respond(message, history, system_message, max_tokens, temperature, top_p):
+    # Build chat prompt
     prompt = system_message + "\n"
     for h in history:
         prompt += f"User: {h['content']}\n"
     prompt += f"User: {message}\nBot:"
 
-    try:
-        # Generate text (non-streaming)
-        output = client.text_generation(
-            prompt,
-            max_new_tokens=max_tokens,
-            temperature=temperature,
-            top_p=top_p,
-        )
-        yield output.text
+    # Tokenize
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+    # Generation config
+    gen_config = GenerationConfig(
+        max_new_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        do_sample=True
+    )
 
-    except Exception as e:
-        yield f"[Error] {str(e)}"
+    # Generate output
+    with torch.no_grad():
+        output_ids = model.generate(**inputs, **gen_config.to_dict())
+    output = tokenizer.decode(output_ids[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
+    return output
 
 # -----------------------------
 # Gradio Chat Interface
@@ -56,8 +69,6 @@ chatbot = gr.ChatInterface(
 )
 
 with gr.Blocks() as demo:
-    with gr.Sidebar():
-        gr.LoginButton()
     chatbot.render()
 
 # -----------------------------