Update app.py

app.py

@@ -4,15 +4,24 @@ import streamlit as st
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.vectorstores import FAISS
 from langchain.embeddings import HuggingFaceEmbeddings
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
+from llama_cpp import Llama
+from huggingface_hub import hf_hub_download  # Add this import
 
 # Configuration
 EMBEDDING_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
-LLM_MODEL_NAME = "TheBloke/Llama-2-7B-GGUF"
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+MODEL_REPO = "TheBloke/Llama-2-7B-GGUF"
+MODEL_FILE = "llama-2-7b.Q4_K_M.gguf"  # Specific quantization file
 
-# Text Splitting
+# Download model from Hugging Face Hub
+@st.cache_resource
+def download_model():
+    return hf_hub_download(
+        repo_id=MODEL_REPO,
+        filename=MODEL_FILE,
+        cache_dir="models"
+    )
+
+# Text Processing
 def get_pdf_text(pdf_docs):
     text = ""
     for pdf in pdf_docs:
@@ -22,135 +31,113 @@ def get_pdf_text(pdf_docs):
     return text
 
 def get_text_chunks(text):
-    splitter = RecursiveCharacterTextSplitter(chunk_size=10000, chunk_overlap=1000)
+    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
     return splitter.split_text(text)
 
 # Vector Store
 def get_vector_store(chunks):
-    try:
-        embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL_NAME)
-        vector_store = FAISS.from_texts(chunks, embedding=embeddings)
-        vector_store.save_local("faiss_index")
-    except Exception as e:
-        raise RuntimeError(f"Error processing embeddings: {e}")
+    embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL_NAME)
+    vector_store = FAISS.from_texts(chunks, embedding=embeddings)
+    vector_store.save_local("faiss_index")
 
-# Load Model with Caching
+# Model Loading
 @st.cache_resource
-def load_model():
-    try:
-        tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL_NAME, cache_dir="./cache", trust_remote_code=True)
-        model = AutoModelForCausalLM.from_pretrained(
-            LLM_MODEL_NAME,
-            torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
-            low_cpu_mem_usage=True,
-            cache_dir="./cache",
-            trust_remote_code=True
-        ).to(DEVICE)
-        return tokenizer, model
-    except Exception as e:
-        raise RuntimeError(f"Error loading DeepSeek model: {e}")
+def load_llama():
+    model_path = download_model()
+    return Llama(
+        model_path=model_path,
+        n_ctx=2048,
+        n_threads=4,
+        n_gpu_layers=20  # Adjust based on your GPU capacity
+    )
 
 # Response Generation
-class DeepSeekChat:
+class LlamaChat:
     def __init__(self):
-        self.tokenizer, self.model = load_model()
+        self.llm = load_llama()
     
     def generate(self, context, question):
-        try:
-            prompt_template = """Answer based on context:
-            Context: {context}
-            Question: {question}
-            Answer:"""
-            
-            prompt = prompt_template.format(context=context, question=question)
-            print("Prompt:", prompt)
-            
-            inputs = self.tokenizer(prompt, return_tensors="pt").to(DEVICE)
-            print("Inputs:", inputs)
-            
-            outputs = self.model.generate(
-                inputs.input_ids,
-                max_new_tokens=200,
-                temperature=0.3,
-                do_sample=True
-            )
-            print("Outputs:", outputs)
-            
-            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-            print("Response:", response)
-            
-            return response
-        except Exception as e:
-            raise RuntimeError(f"Error generating response: {e}")
-
-# Streamlit App
-def clear_chat_history():
-    st.session_state.messages = [{"role": "assistant", "content": "در خدمتیم"}]
+        prompt = f"""Answer based on this context:
+        {context}
+        
+        Question: {question}
+        Answer:"""
+        
+        response = self.llm(
+            prompt,
+            max_tokens=256,
+            temperature=0.3,
+            top_p=0.95,
+            repeat_penalty=1.2
+        )
+        
+        return response['choices'][0]['text'].strip()
 
+# Streamlit UI
 def main():
-    st.set_page_config(page_title="Chatbot", layout="wide", initial_sidebar_state="expanded")
+    st.set_page_config(page_title="Llama-2 PDF Assistant", layout="wide")
     
-    st.markdown("""
-    <style>
-    body { background-color: #000000; color: #ffffff; }
-    .main { background-color: #333333; padding: 20px; border-radius: 10px; }
-    </style>
-    """, unsafe_allow_html=True)
-
+    # Initialize session state
     if "uploaded" not in st.session_state:
         st.session_state.uploaded = False
+    if "messages" not in st.session_state:
+        st.session_state.messages = [{"role": "assistant", "content": "Ask me about your documents!"}]
 
+    # File Upload Section
     if not st.session_state.uploaded:
-        st.title("A Personal Assistant for your Standards ...")
-        pdf_docs = st.file_uploader("فایل مورد نظر راآپلود کنید", accept_multiple_files=True)
-        if st.button("تایید"):
+        st.title("PDF Document Assistant")
+        pdf_docs = st.file_uploader("Upload PDF files", accept_multiple_files=True, type="pdf")
+        
+        if st.button("Process Documents"):
             if pdf_docs:
-                with st.spinner("در حال پردازش..."):
+                with st.status("Processing documents..."):
+                    st.write("Extracting text...")
                     raw_text = get_pdf_text(pdf_docs)
+                    
+                    st.write("Chunking text...")
                     text_chunks = get_text_chunks(raw_text)
+                    
+                    st.write("Creating vector store...")
                     get_vector_store(text_chunks)
+                    
                     st.session_state.uploaded = True
-                    st.success("پردازش موفق")
+                    st.rerun()
             else:
-                st.error("لطفا حداقل یک فایل را آپلود کنید")
-    else:
-        st.title("Assistant ready ...")
-        st.write("میتونین سوالتونو بپرسین")
-        st.button('حذف مکالمه', on_click=clear_chat_history)
-
-        if "messages" not in st.session_state:
-            st.session_state.messages = [{"role": "assistant", "content": "در خدمتیم"}]
-
-        for message in st.session_state.messages:
-            with st.chat_message(message["role"]):
-                st.write(message["content"])
+                st.error("Please upload at least one PDF file")
 
-        if prompt := st.chat_input():
+    # Chat Interface
+    else:
+        st.title("Document Chat Assistant")
+        st.button("Clear Chat", on_click=lambda: st.session_state.update(
+            messages=[{"role": "assistant", "content": "Ask me about your documents!"}]
+        ))
+        
+        # Display chat messages
+        for msg in st.session_state.messages:
+            st.chat_message(msg["role"]).write(msg["content"])
+        
+        # Handle user input
+        if prompt := st.chat_input("Your question:"):
             st.session_state.messages.append({"role": "user", "content": prompt})
-            with st.chat_message("user"):
-                st.write(prompt)
-
-            if st.session_state.messages[-1]["role"] != "assistant":
-                with st.spinner("در حال پاسخگویی..."):
-                    try:
-                        embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL_NAME)
-                        db = FAISS.load_local("faiss_index", embeddings, allow_dangerous_deserialization=True)
-                        docs = db.similarity_search(prompt, k=4)
-                        print("Retrieved Documents:", docs)
-                        
-                        context = "\n".join([d.page_content for d in docs])
-                        print("Context:", context)
-                        
-                        chat_model = DeepSeekChat()
-                        response = chat_model.generate(context, prompt)
-                        print("Generated Response:", response)
-                        
-                        with st.chat_message("assistant"):
-                            st.write(response)
-                            st.session_state.messages.append({"role": "assistant", "content": response})
-                    except Exception as e:
-                        st.error(f"خطا: {str(e)}")
-                        print("Error during response generation:", str(e))
+            st.chat_message("user").write(prompt)
+            
+            with st.spinner("Thinking..."):
+                try:
+                    # Retrieve relevant context
+                    embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL_NAME)
+                    db = FAISS.load_local("faiss_index", embeddings, allow_dangerous_deserialization=True)
+                    docs = db.similarity_search(prompt, k=4)
+                    context = "\n".join([d.page_content for d in docs])
+                    
+                    # Generate response
+                    chat_model = LlamaChat()
+                    response = chat_model.generate(context, prompt)
+                    
+                    st.session_state.messages.append({"role": "assistant", "content": response})
+                    st.chat_message("assistant").write(response)
+                    
+                except Exception as e:
+                    st.error(f"Error: {str(e)}")
 
 if __name__ == "__main__":
     main()