Update app.py

app.py

@@ -3,7 +3,7 @@ import os
 import re
 import logging
 import torch
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from playwright.sync_api import sync_playwright
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import FAISS
@@ -14,13 +14,12 @@ logging.basicConfig(
     filename='/app/cache/app.log',
     level=logging.DEBUG,
     format='%(asctime)s - %(levelname)s - %(message)s'
-
 )
 
-MODEL_NAME = "google/flan-t5-large"
+MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 
 st.set_page_config(
-    page_title="RAG · Mistral",
+    page_title="RAG · TinyLlama",
     page_icon="🕸️",
     layout="wide",
     initial_sidebar_state="collapsed"
@@ -29,7 +28,6 @@ st.set_page_config(
 st.markdown("""
 <style>
 @import url('https://fonts.googleapis.com/css2?family=Instrument+Serif:ital@0;1&family=JetBrains+Mono:wght@300;400;500&display=swap');
-
 :root {
     --bg:      #f5f0e8;
     --surface: #ede8df;
@@ -40,7 +38,6 @@ st.markdown("""
     --mono:    'JetBrains Mono', monospace;
     --serif:   'Instrument Serif', serif;
 }
-
 html, body, [class*="css"] {
     font-family: var(--mono);
     background: var(--bg);
@@ -49,12 +46,10 @@ html, body, [class*="css"] {
 .stApp { background: var(--bg); }
 #MainMenu, footer, header { visibility: hidden; }
 [data-testid="stDecoration"] { display: none; }
-
 [data-testid="stSidebar"] {
     background: var(--surface);
     border-right: 1px solid var(--border);
 }
-
 .stTextInput > div > div > input,
 .stTextArea textarea {
     background: #fff !important;
@@ -69,7 +64,6 @@ html, body, [class*="css"] {
     border-color: var(--accent) !important;
     box-shadow: 0 0 0 2px rgba(193,58,30,0.12) !important;
 }
-
 .stButton > button {
     background: var(--accent) !important;
     color: #fff !important;
@@ -88,7 +82,6 @@ html, body, [class*="css"] {
     transform: translateY(-1px);
     box-shadow: 0 3px 12px rgba(193,58,30,0.25) !important;
 }
-
 [data-testid="stChatMessage"] {
     background: #fff !important;
     border: 1px solid var(--border) !important;
@@ -100,9 +93,7 @@ html, body, [class*="css"] {
     font-family: var(--mono) !important;
     font-size: 0.82rem !important;
 }
-
 hr { border-color: var(--border) !important; }
-
 .content-box {
     background: #fff;
     border: 1px solid var(--border);
@@ -120,7 +111,6 @@ hr { border-color: var(--border) !important; }
 .content-box::-webkit-scrollbar { width: 6px; }
 .content-box::-webkit-scrollbar-track { background: var(--surface); }
 .content-box::-webkit-scrollbar-thumb { background: var(--border); border-radius: 3px; }
-
 .meta-pill {
     display: inline-flex;
     align-items: center;
@@ -134,7 +124,6 @@ hr { border-color: var(--border) !important; }
     margin-bottom: 0.6rem;
 }
 .meta-dot { width:6px; height:6px; border-radius:50%; background:#4caf50; }
-
 .section-label {
     font-size: 0.68rem;
     letter-spacing: 0.12em;
@@ -151,7 +140,6 @@ hr { border-color: var(--border) !important; }
     height: 1px;
     background: var(--border);
 }
-
 .qa-banner {
     display: flex;
     align-items: center;
@@ -166,8 +154,7 @@ hr { border-color: var(--border) !important; }
     color: var(--accent);
     white-space: nowrap;
 }
-
-.ollama-badge {
+.model-badge {
     display: inline-flex;
     align-items: center;
     gap: 5px;
@@ -177,8 +164,7 @@ hr { border-color: var(--border) !important; }
     border: 1px solid var(--border);
     border-radius: 3px;
 }
-.ollama-dot { width:6px; height:6px; border-radius:50%; }
-
+.model-dot { width:6px; height:6px; border-radius:50%; }
 .page-header {
     padding: 1.5rem 0 1rem 0;
     border-bottom: 2px solid var(--text);
@@ -202,7 +188,6 @@ hr { border-color: var(--border) !important; }
     letter-spacing: 0.08em;
     text-transform: uppercase;
 }
-
 [data-testid="stAlert"] {
     background: var(--surface) !important;
     border: 1px solid var(--border) !important;
@@ -228,7 +213,7 @@ for key, default in [
 # ── Utilities ──────────────────────────────────────────────────────────────────
 
 def clean_text(text):
-    # Only collapse whitespace — preserve Rs. prices, commas, symbols
+    # Only collapse whitespace — preserve prices, commas, symbols
     text = re.sub(r'[ \t]+', ' ', text)
     text = re.sub(r'\n{3,}', '\n\n', text)
     return text.strip()
@@ -236,14 +221,17 @@ def clean_text(text):
 def is_valid_url(url):
     return bool(re.match(r'^https?://[\w\-\.]+(?::\d+)?(?:/[\w\-\./]*)*$', url))
 
-def check_model():
-    return st.session_state.get('qa_model') is not None
+# ── Model ──────────────────────────────────────────────────────────────────────
 
 @st.cache_resource(show_spinner=False)
 def load_model():
     try:
         tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-        model     = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float32)
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_NAME,
+            torch_dtype=torch.float32,
+            low_cpu_mem_usage=True,
+        )
         model.eval()
         logging.info(f"Loaded {MODEL_NAME}")
         return tokenizer, model
@@ -251,50 +239,47 @@ def load_model():
         logging.error(f"Model load error: {e}")
         return None, None
 
+# ── Scraper ────────────────────────────────────────────────────────────────────
+
 def scrape_website(url):
     with sync_playwright() as p:
-        browser = p.chromium.launch(headless=True, args=['--no-sandbox','--disable-dev-shm-usage'])
+        browser = p.chromium.launch(headless=True, args=['--no-sandbox', '--disable-dev-shm-usage'])
         page = browser.new_page()
         try:
             page.goto(url, wait_until="networkidle", timeout=45000)
             title = page.title()
 
-            # Strategy 1: extract structured name+price pairs from <li> elements
-            # Works well for listing/price pages like whatmobile.com.pk
+            # Strategy 1: extract from <li> elements — good for listing/price pages
             lines = []
-            li_elements = page.query_selector_all("li")
-            for li in li_elements:
+            for li in page.query_selector_all("li"):
                 try:
                     text = li.inner_text().strip()
-                    # Keep li items that contain a heading and a price-like pattern
-                    if text and len(text) > 3 and len(text) < 300:
+                    if text and 3 < len(text) < 300:
                         lines.append(text)
                 except:
                     continue
 
-            # Strategy 2: grab all headings and paragraphs too
+            # Strategy 2: headings, paragraphs, table cells
             for tag in ["h1", "h2", "h3", "h4", "p", "td"]:
-                elements = page.query_selector_all(tag)
-                for e in elements:
+                for e in page.query_selector_all(tag):
                     try:
                         text = e.inner_text().strip()
-                        if text and len(text) > 3 and len(text) < 500:
+                        if text and 3 < len(text) < 500:
                             lines.append(text)
                     except:
                         continue
 
-            # Deduplicate while preserving order
-            seen = set()
-            unique_lines = []
+            # Deduplicate preserving order
+            seen, unique_lines = set(), []
             for line in lines:
-                normalised = re.sub(r'\s+', ' ', line).strip()
-                if normalised not in seen:
-                    seen.add(normalised)
-                    unique_lines.append(normalised)
+                n = re.sub(r'\s+', ' ', line).strip()
+                if n not in seen:
+                    seen.add(n)
+                    unique_lines.append(n)
 
             content = "\n".join(unique_lines)
 
-            # Fallback to full body if we got almost nothing
+            # Fallback to body if nothing found
             if len(content) < 200:
                 body = page.query_selector("body")
                 content = clean_text(body.inner_text()) if body else content
@@ -308,6 +293,8 @@ def scrape_website(url):
         finally:
             browser.close()
 
+# ── Vector store ───────────────────────────────────────────────────────────────
+
 @st.cache_resource
 def create_vector_store(text):
     try:
@@ -322,6 +309,8 @@ def create_vector_store(text):
         st.error(f"Indexing failed: {e}")
         return None
 
+# ── Answer ─────────────────────────────────────────────────────────────────────
+
 def answer_question(question):
     if not st.session_state.vector_store:
         return "No content indexed yet."
@@ -329,23 +318,55 @@ def answer_question(question):
     if tokenizer is None:
         return "Model failed to load. Check logs."
     try:
+        # Retrieve top 3 relevant chunks from FAISS
         docs    = st.session_state.vector_store.similarity_search(question, k=3)
         context = " ".join(d.page_content for d in docs)
-        prompt  = (
-            "Answer the question using only the context provided. "
-            "If the answer is not in the context, say \"I don't know\".\n\n"
-            f"Context: {context}\n\nQuestion: {question}\n\nAnswer:"
+
+        # TinyLlama expects the chat template format
+        messages = [
+            {
+                "role": "system",
+                "content": (
+                    "You are a helpful assistant. Answer the user's question using "
+                    "ONLY the context provided. If the answer is not in the context, "
+                    "say \"I don't know\"."
+                ),
+            },
+            {
+                "role": "user",
+                "content": f"Context:\n{context}\n\nQuestion: {question}",
+            },
+        ]
+
+        # Apply chat template → produces <|system|>...<|user|>...<|assistant|>
+        prompt = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True,   # appends <|assistant|> so model starts answering
+        )
+
+        inputs = tokenizer(
+            prompt,
+            return_tensors="pt",
+            truncation=True,
+            max_length=2048,             # TinyLlama's full context window
         )
-        inputs  = tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True)
+
         with torch.no_grad():
             outputs = model.generate(
                 **inputs,
-                max_new_tokens=200,
-                num_beams=4,
-                early_stopping=True,
-                no_repeat_ngram_size=3,
+                max_new_tokens=300,
+                do_sample=True,
+                temperature=0.7,
+                top_p=0.95,
+                repetition_penalty=1.1,
+                pad_token_id=tokenizer.eos_token_id,
             )
-        return tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
+
+        # Slice off the prompt tokens — only decode what the model generated
+        generated = outputs[0][inputs["input_ids"].shape[1]:]
+        return tokenizer.decode(generated, skip_special_tokens=True).strip()
+
     except Exception as e:
         logging.error(f"Inference error: {e}")
         return f"Error generating answer: {e}"
@@ -359,6 +380,8 @@ model_ok = _tok is not None
 with st.sidebar:
     st.markdown("**Model**")
     st.markdown(f"`{MODEL_NAME}`")
+    st.markdown("**Context window**")
+    st.markdown("`2048 tokens`")
     st.markdown("**Status**")
     if model_ok:
         st.success("Model loaded ✓")
@@ -375,9 +398,9 @@ st.markdown(f"""
         <p class="page-title">Web RAG</p>
         <span class="page-sub">Scrape → Index → Ask</span>
     </div>
-    <div class="ollama-badge">
-        <div class="ollama-dot" style="background:{dot_color};"></div>
-        {dot_label} &nbsp;·&nbsp; flan-t5-large
+    <div class="model-badge">
+        <div class="model-dot" style="background:{dot_color};"></div>
+        {dot_label} &nbsp;·&nbsp; TinyLlama-1.1B-Chat
     </div>
 </div>
 """, unsafe_allow_html=True)
@@ -411,7 +434,6 @@ if scrape_clicked:
 # ── Main content area ──────────────────────────────────────────────────────────
 if st.session_state.scraped_content:
 
-    # Meta pill
     title_display = st.session_state.scraped_title or ""
     url_display   = st.session_state.scraped_url or ""
     st.markdown(f"""
@@ -421,18 +443,16 @@ if st.session_state.scraped_content:
         &nbsp;·&nbsp;
         <span>{st.session_state.char_count:,} chars</span>
         &nbsp;·&nbsp;
-        <span style="max-width:300px; overflow:hidden; text-overflow:ellipsis; white-space:nowrap;">{url_display}</span>
+        <span style="max-width:300px;overflow:hidden;text-overflow:ellipsis;white-space:nowrap;">{url_display}</span>
     </div>
     """, unsafe_allow_html=True)
 
-    # Scraped content label + scrollable box
     st.markdown('<div class="section-label">Scraped content</div>', unsafe_allow_html=True)
     preview = st.session_state.scraped_content[:4000]
     if len(st.session_state.scraped_content) > 4000:
         preview += "\n\n… (truncated for display)"
     st.markdown(f'<div class="content-box">{preview}</div>', unsafe_allow_html=True)
 
-    # ── Q&A section directly below ─────────────────────────────────────────────
     st.markdown("""
     <div class="qa-banner">
         <div class="qa-banner-line"></div>
@@ -441,18 +461,16 @@ if st.session_state.scraped_content:
     </div>
     """, unsafe_allow_html=True)
 
-    # Render past exchanges
     for msg in st.session_state.chat_history:
         with st.chat_message(msg["role"]):
             st.markdown(msg["content"])
 
-    # Chat input
     if prompt := st.chat_input("Ask anything about the content above…"):
         st.session_state.chat_history.append({"role": "user", "content": prompt})
         with st.chat_message("user"):
             st.markdown(prompt)
         with st.chat_message("assistant"):
-            with st.spinner("FLAN-T5 is thinking…"):
+            with st.spinner("TinyLlama is thinking…"):
                 answer = answer_question(prompt)
             st.markdown(answer)
         st.session_state.chat_history.append({"role": "assistant", "content": answer})
@@ -463,7 +481,6 @@ if st.session_state.scraped_content:
             st.rerun()
 
 else:
-    # Empty state
     st.markdown("""
     <div style="
         text-align:center;