Update app.py

app.py

@@ -1,75 +1,35 @@
-import os
-import torch
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from peft import PeftModel, PeftConfig
+from llama_cpp import Llama
 
-# Set the HF repo and LoRA model location
-base_model_id = "unsloth/gemma-2-9b"
-lora_model_id = "Futuresony/gemma2-9b-lora-alpaca"
+# Path to your GGUF model inside the space
+MODEL_PATH = "your-model.gguf"
 
-# Load base model on CPU
-base_model = AutoModelForCausalLM.from_pretrained(
-    base_model_id,
-    device_map="cpu",
-    torch_dtype=torch.float32,
-)
+# Load model
+llm = Llama(model_path=MODEL_PATH, n_ctx=2048, n_threads=4, verbose=True)
 
-# Load tokenizer from base model
-tokenizer = AutoTokenizer.from_pretrained(base_model_id)
-
-# Load LoRA adapter
-model = PeftModel.from_pretrained(base_model, lora_model_id)
-model.eval()
-
-# === Alpaca-style formatter ===
-def format_alpaca_prompt(user_input, system_prompt, history):
-    history_str = "\n".join([f"### Instruction:\n{h[0]}\n### Response:\n{h[1]}" for h in history])
-    prompt = f"""{system_prompt}
-{history_str}
-
-### Instruction:
-{user_input}
+# Function to format the prompt
+def format_prompt(user_message):
+    return f"""### Instruction:
+{user_message}
 
 ### Response:"""
-    return prompt
-
-# === Chat logic ===
-def respond(message, history, system_message, max_tokens, temperature, top_p):
-    prompt = format_alpaca_prompt(message, system_message, history)
-    inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
 
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=max_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id,
-        )
+# Chat handler
+def respond(user_message, chat_history):
+    prompt = format_prompt(user_message)
+    output = llm(prompt, max_tokens=300, stop=["###"])
+    response = output["choices"][0]["text"].strip()
+    chat_history.append((user_message, response))
+    return "", chat_history
 
-    response_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    # Only return the part after "### Response:"
-    if "### Response:" in response_text:
-        final_output = response_text.split("### Response:")[-1].strip()
-    else:
-        final_output = response_text.strip()
+# Gradio UI
+with gr.Blocks() as demo:
+    gr.Markdown("## 🤖 DStv AI Assistant (Offline - GGUF)")
+    chatbot = gr.Chatbot()
+    msg = gr.Textbox(placeholder="Ask your question...")
+    state = gr.State([])
 
-    history.append((message, final_output))
-    yield final_output
+    msg.submit(respond, [msg, state], [msg, chatbot])
 
-# === Gradio Interface ===
-demo = gr.ChatInterface(
-    fn=respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=1024, value=256, 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.01, label="Top-p"),
-    ],
-    title="Offline Gemma-2B Alpaca Chatbot (LoRA)",
-)
+demo.launch()
 
-if __name__ == "__main__":
-    demo.launch()