Upload app.py

app.py

@@ -2,10 +2,39 @@ import gradio as gr
 from llama_cpp import Llama
 from huggingface_hub import hf_hub_download
 import os
+import pickle
+from langchain_huggingface import HuggingFaceEmbeddings
 
-# Download model
+# ====================== CONFIG ======================
+repo_id = "robertolofaro/books-model"
+
+BACKENDS = {
+    "Fast Mode (No RAG)": None,
+    "Chroma - RAG": "Chroma",
+    "FAISS - RAG (HNSW)": "FAISS",
+    "Qdrant - RAG": "Qdrant"
+}
+
+CHROMA_PATH = "chroma_db"
+FAISS_PATH = "faiss_index_hnsw"
+QDRANT_PATH = "qdrant_db"
+QDRANT_COLLECTION = "articles"
+
+# ====================== LOAD METADATA FOR BOOK LIST ======================
+def load_articles_list():
+    try:
+        with open("metadata.pkl", "rb") as f:
+            df = pickle.load(f)
+        articles = sorted(df['article_category'].unique().tolist())
+        return ["All categories"] + articles
+    except:
+        return ["All categories"]
+
+ARTICLE_LIST = load_article_list()
+
+# ====================== LOAD LLM ======================
 model_path = hf_hub_download(
-    repo_id="robertolofaro/articles-model",
+    repo_id=repo_id,
     filename="articles-Q4_K_M.gguf",
     repo_type="model",
     token=os.environ.get("HF_TOKEN")
@@ -13,50 +42,128 @@ model_path = hf_hub_download(
 
 llm = Llama(
     model_path=model_path,
-    n_ctx=65000,
+    n_ctx=4096,
     n_threads=2,
     n_batch=512,
     n_ubatch=512,
-    n_gpu_layers=0,
-    verbose=False
+    verbose=False,
 )
 
+# ====================== RAG CACHE ======================
+vectorstores = {}
+
+def get_vectorstore(backend_name: str):
+    if backend_name in vectorstores:
+        return vectorstores[backend_name]
+    # ... (same loading logic as before - Chroma, FAISS, Qdrant) ...
+    # I'll keep it short here for brevity, but same as previous version
+    try:
+        embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-small-en-v1.5", encode_kwargs={'normalize_embeddings': True})
+        
+        if backend_name == "Chroma":
+            from langchain_community.vectorstores import Chroma
+            vs = Chroma(persist_directory=CHROMA_PATH, embedding_function=embeddings)
+        elif backend_name == "FAISS":
+            from langchain_community.vectorstores import FAISS
+            vs = FAISS.load_local(FAISS_PATH, embeddings, allow_dangerous_deserialization=True)
+        elif backend_name == "Qdrant":
+            from langchain_community.vectorstores import Qdrant
+            vs = Qdrant(path=QDRANT_PATH, collection_name=QDRANT_COLLECTION, embeddings=embeddings)
+        else:
+            return None
+            
+        vectorstores[backend_name] = vs
+        return vs
+    except:
+        return None
+
+
+# ====================== SYSTEM PROMPT ======================
 SYSTEM_PROMPT = """You are the reference expert for the articles contained in this database, all extracted from the website robertolofaro.com, and all focused on change.
 #Your Mission:
-When a user asks a question, your goal is to provide a structured response based ONLY on the articles provided in your training. Do not provide general advice from outside these sources. Do not reference specific articles or their titles, only the argument that they represent
+When a user asks a question, your goal is to provide a structured response based ONLY on the articles provided in your training. Do not provide general advice from outside these sources.
 # Response Format:
 1. Executive Summary: A 2-3 sentence overview answering the core query.
-2. Guidelines & Hints: A narrative of specific "answers/guidelines/hints" found in the source material, but without repetition and as storytelling. Verify and summarize your answer before sharing it with the user, so that it is not verbose and as short as possible- no more than 250 words should be within your answer.
+2. Guidelines & Hints: A markdown list of specific "answers/guidelines/hints" found in the source material.
 """
 
-def generate_response(message, history):
+
+# ====================== GENERATION FUNCTION ======================
+def generate_response(message, history, rag_mode, book_filter, max_tokens, temperature, top_p, repeat_penalty):
     full_prompt = f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n"
-    
-    for msg in history:
-        role = msg["role"]
-        content = msg["content"]
-        full_prompt += f"<|im_start|>{role}\n{content}<|im_end|>\n"
-    
-    full_prompt += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
+
+    for msg in history[-4:]:
+        full_prompt += f"<|im_start|>{msg['role']}\n{msg['content']}<|im_end|>\n"
+
+    backend = BACKENDS.get(rag_mode)
+    context = ""
+
+    if backend:
+        vs = get_vectorstore(backend)
+        if vs:
+            try:
+                filter_dict = {"article_category": article_filter} if article_filter != "All categories" else None
+                docs = vs.similarity_search(message, k=5, filter=filter_dict)
+                context = "\n\n".join([
+                    f"[Category: {doc.metadata.get('article_category', 'N/A')}] {doc.page_content[:700]}"
+                    for doc in docs
+                ])
+            except:
+                pass
+
+    if context:
+        full_prompt += f"<|im_start|>user\nContext:\n{context}\n\nQuestion: {message}<|im_end|>\n"
+    else:
+        full_prompt += f"<|im_start|>user\n{message}<|im_end|>\n"
+
+    full_prompt += "<|im_start|>assistant\n"
 
     partial_text = ""
     for chunk in llm(
         full_prompt,
-        max_tokens=2048,
+        max_tokens=int(max_tokens),
+        temperature=float(temperature),
+        top_p=float(top_p),
+        repeat_penalty=float(repeat_penalty),
         stop=["<|im_end|>", "<|im_start|>"],
         stream=True,
-        temperature=0.7,
     ):
         token = chunk['choices'][0]['text']
         partial_text += token
         yield partial_text
 
 
-demo = gr.ChatInterface(
-    fn=generate_response,
-    title="Articles Q&A (CPU)",
-    description="Experimental Q&A based on 350+ articles on Change Management published on robertolofaro.com. It could take few minutes to start delivering an answer (CPU-based).",
-)
+# ====================== GRADIO INTERFACE ======================
+with gr.Blocks(title="Article Q&A model") as demo:
+    gr.Markdown("# sourcing 350+ articles on change")
+    gr.Markdown("Qwen3.5-4B DoRA fine-tuned on 350+ articles")
+
+    with gr.Row():
+        rag_mode = gr.Radio(
+            choices=list(BACKENDS.keys()),
+            value="Fast Mode (No RAG)",
+            label="Mode"
+        )
+        book_filter = gr.Dropdown(
+            choices=BOOK_LIST,
+            value="All categories",
+            label="Focus on category"
+        )
+
+    with gr.Accordion("Advanced Generation Parameters", open=False):
+        max_tokens = gr.Slider(256, 2048, value=900, step=64, label="Max Tokens")
+        temperature = gr.Slider(0.0, 1.0, value=0.65, step=0.05, label="Temperature")
+        top_p = gr.Slider(0.0, 1.0, value=0.9, step=0.05, label="Top-p")
+        repeat_penalty = gr.Slider(1.0, 2.0, value=1.1, step=0.05, label="Repeat Penalty")
+
+    gr.ChatInterface(
+        fn=generate_response,
+        additional_inputs=[rag_mode, book_filter, max_tokens, temperature, top_p, repeat_penalty],
+        examples=[
+            ["What is the potential for Italy?"],
+            ["What is the potential for Turin?"]
+        ],
+    )
 
 if __name__ == "__main__":
     demo.queue(default_concurrency_limit=1).launch()