updated

app.py

@@ -1,10 +1,8 @@
-# app.py — CPU-ready with Gemma 2B
 import gradio as gr
 import os
 import tempfile
 import textwrap
 from datetime import datetime
-from pathlib import Path
 from typing import List, Dict, Any, Optional
 
 from src.model_loader import load_local_model
@@ -14,7 +12,7 @@ from src.chatbot import LocalChatbot
 # ----------------------
 # Model setup
 # ----------------------
-MODEL_PATH = "models/gemma-2-2b-it-Q4_K_M.gguf"  # quantized 2B
+MODEL_PATH = "models/gemma-2-2b-it-Q4_K_M"  # quantized 2B
 llm = load_local_model(MODEL_PATH, device=-1)  # CPU
 memory = ConversationMemory(max_len=60)
 bot = LocalChatbot(llm, memory)
@@ -103,7 +101,7 @@ def generate_reply(user_msg: str, history: Optional[List[Dict[str, Any]]]):
     return history
 
 # ----------------------
-# UI
+# Gradio UI
 # ----------------------
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     with gr.Row():

src/model_loader.py

@@ -1,13 +1,18 @@
-# src/model_loader.py
-from llama_cpp import Llama
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
-def load_local_model(model_path):
-    print(f"Loading model: {model_path}")
-    llm = Llama(
-        model_path=model_path,
-        n_ctx=4096,
-        n_threads=6,
-        n_gpu_layers=0,
-        verbose=False
+def load_local_model(model_path, device=0):
+    """
+    Loads a local quantized model for CPU or GPU.
+    device=-1 => CPU, device>=0 => GPU
+    """
+    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    model = AutoModelForCausalLM.from_pretrained(model_path)
+    
+    # Use pipeline for text generation
+    generator = pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        device=device
     )
-    return llm
+    return generator