Update app.py

app.py

@@ -3,62 +3,53 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 from peft import PeftModel
 import torch
 
-# Define the base and LoRA model IDs
-base_model_id = "unsloth/gemma-2-9b-bnb-4bit"
+# Use a CPU-compatible base model (replace this with your actual full-precision model)
+base_model_id = "unsloth/gemma-2b"  # Replace with real CPU-compatible model
 lora_model_id = "Futuresony/future_12_10_2024"
 
-# Load the base model on CPU with float16
+# Load the base model on CPU
 base_model = AutoModelForCausalLM.from_pretrained(
     base_model_id,
-    torch_dtype=torch.float16,
-    device_map="cpu",  # Load the model on CPU, no GPU
+    torch_dtype=torch.float32,  # Use float32 for CPU compatibility
+    device_map="cpu"
 )
 
 # Load the PEFT LoRA model
 model = PeftModel.from_pretrained(base_model, lora_model_id)
 
-# Tokenizer for the model
+# Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(base_model_id)
 
-# Function to respond to the user's input
+# Chat function
 def respond(message, history: list[tuple[str, str]], system_message, max_tokens, temperature, top_p):
-    # Prepare the message history for chat completion
     messages = [{"role": "system", "content": system_message}]
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-
+    for user_msg, bot_msg in history:
+        if user_msg:
+            messages.append({"role": "user", "content": user_msg})
+        if bot_msg:
+            messages.append({"role": "assistant", "content": bot_msg})
     messages.append({"role": "user", "content": message})
 
-    # Generate a response
-    response = ""
-    for message in model.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
+    # Generate response (simulated loop for streaming output)
+    inputs = tokenizer.apply_chat_template(messages, return_tensors="pt").to("cpu")
+    outputs = model.generate(
+        inputs,
+        max_new_tokens=max_tokens,
         temperature=temperature,
         top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
-        response += token
-        yield response
+        do_sample=True,
+    )
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    yield response
 
-# Gradio interface setup
+# Gradio UI
 demo = gr.ChatInterface(
-    respond,
+    fn=respond,
     additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
+        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)",
-        ),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p"),
     ],
 )