Update app.py

app.py

@@ -1,19 +1,21 @@
 import os
 from PyPDF2 import PdfReader
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.embeddings import HuggingFaceEmbeddings  # Corrected import
 import streamlit as st
+from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.vectorstores import FAISS
-from langchain.chains.question_answering import load_qa_chain
 from langchain.prompts import PromptTemplate
-from dotenv import load_dotenv
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from sentence_transformers import SentenceTransformer
 import torch
 
-# Load environment variables
-load_dotenv()
+# ========== Configuration ==========
+EMBEDDING_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"  # Replace with actual model name
+LLM_MODEL_NAME = "deepseek-ai/Janus-Pro-7B"  # Replace with actual model name
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
 
-# Function to read PDF files
+
+# ========== Text Splitting ==========
 def get_pdf_text(pdf_docs):
     text = ""
     for pdf in pdf_docs:
@@ -22,13 +24,11 @@ def get_pdf_text(pdf_docs):
             text += page.extract_text()
     return text
 
-# Function to split text into chunks
 def get_text_chunks(text):
     splitter = RecursiveCharacterTextSplitter(chunk_size=10000, chunk_overlap=1000)
-    chunks = splitter.split_text(text)
-    return chunks
+    return splitter.split_text(text)
 
-# Function to generate and save embeddings
+# ========== Vector Store ==========
 def get_vector_store(chunks):
     try:
         embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
@@ -37,124 +37,101 @@ def get_vector_store(chunks):
     except Exception as e:
         raise RuntimeError(f"Error processing embeddings: {e}")
 
-# Function to get conversational chain
-def get_conversational_chain():
-    prompt_template = """
-    Answer the question as detailed as possible from the provided context. 
-    If the answer is not in the provided context, just say, "answer is not available in the context".
-    Don't provide the wrong answer.\n\n
-    Context:\n {context}\n
-    Question:\n {question}\n
-    Answer:
-    """
-    try:
-        model = pipeline("text-generation", model="EleutherAI/gpt-j-6B", device="cuda" if torch.cuda.is_available() else "cpu")
-        prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
+# ========== Response Generation ==========
+class DeepSeekChat:
+    def __init__(self):
+        self.tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL_NAME)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            LLM_MODEL_NAME,
+            torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
+            low_cpu_mem_usage=True
+        ).to(DEVICE)
+    
+    def generate(self, context, question):
+        prompt_template = """Answer based on context:
+        Context: {context}
+        Question: {question}
+        Answer:"""
         
-        def generate_response(context, question):
-            input_text = prompt.format(context=context, question=question)
-            response = model(input_text, max_length=200, num_return_sequences=1)
-            return response[0]['generated_text']
+        prompt = prompt_template.format(context=context, question=question)
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(DEVICE)
         
-        return generate_response
-    except Exception as e:
-        raise RuntimeError(f"Error creating conversational chain: {e}")
+        outputs = self.model.generate(
+            inputs.input_ids,
+            max_new_tokens=200,
+            temperature=0.3,
+            do_sample=True
+        )
+        
+        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-# Function to clear chat history
+# ========== Streamlit App ==========
 def clear_chat_history():
     st.session_state.messages = [{"role": "assistant", "content": "در خدمتیم"}]
 
-# Function to handle user input
-def user_input(user_question):
-    try:
-        embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")  # Corrected here
-        new_db = FAISS.load_local("faiss_index", embeddings, allow_dangerous_deserialization=True)
-        docs = new_db.similarity_search(user_question, k=4)
-        chain = get_conversational_chain()
-        response = chain({"input_documents": docs, "question": user_question}, return_only_outputs=True)
-        return response
-    except Exception as e:
-        raise RuntimeError(f"Error generating response: {e}")
-
-# Main function to run the Streamlit app
 def main():
-    st.set_page_config(
-        page_title="Chatbot",
-        layout="wide",
-        initial_sidebar_state="expanded"
-    )
-
-    # Custom CSS for better design
-    st.markdown(
-        """
-        <style>
-        body {
-            background-color: #000000;
-            color: #ffffff;
-        }
-        .main {
-            background-color: #333333;
-            padding: 20px;
-            border-radius: 10px;
-        }
-        </style>
-        """,
-        unsafe_allow_html=True
-    )
+    st.set_page_config(page_title="Chatbot", layout="wide", initial_sidebar_state="expanded")
+    
+    # Custom CSS
+    st.markdown("""
+    <style>
+    body { background-color: #000000; color: #ffffff; }
+    .main { background-color: #333333; padding: 20px; border-radius: 10px; }
+    </style>
+    """, unsafe_allow_html=True)
 
     if "uploaded" not in st.session_state:
         st.session_state.uploaded = False
 
     if not st.session_state.uploaded:
         st.title("A Personal Assistant for your Standards ...")
-        pdf_docs = st.file_uploader("فایل مورد نظر راآپلود کنید ", accept_multiple_files=True)
+        pdf_docs = st.file_uploader("فایل مورد نظر راآپلود کنید", accept_multiple_files=True)
         if st.button("تایید"):
             if pdf_docs:
-                try:
-                    st.info("در حال پردازش")
+                with st.spinner("در حال پردازش..."):
                     raw_text = get_pdf_text(pdf_docs)
                     text_chunks = get_text_chunks(raw_text)
                     get_vector_store(text_chunks)
                     st.session_state.uploaded = True
                     st.success("پردازش موفق")
-                    st.query_params["uploaded"] = "true"
-                except RuntimeError as e:
-                    st.error(str(e))
             else:
                 st.error("لطفا حداقل یک فایل را آپلود کنید")
     else:
-        # Main content area for displaying chat messages
         st.title("Assistant ready ...")
         st.write("میتونین سوالتونو بپرسین")
         st.button('حذف مکالمه', on_click=clear_chat_history)
 
-        # Initialize chat history
         if "messages" not in st.session_state:
             st.session_state.messages = [{"role": "assistant", "content": "در خدمتیم"}]
 
-        # Display chat messages
         for message in st.session_state.messages:
             with st.chat_message(message["role"]):
                 st.write(message["content"])
 
-        # Chat input
         if prompt := st.chat_input():
             st.session_state.messages.append({"role": "user", "content": prompt})
             with st.chat_message("user"):
                 st.write(prompt)
 
-            # Generate bot response
             if st.session_state.messages[-1]["role"] != "assistant":
-                try:
-                    with st.chat_message("assistant"):
-                        response = user_input(prompt)
-                        if response:
-                            full_response = ''.join(response['output_text'])
-                            st.write(full_response)
-                            message = {"role": "assistant", "content": full_response}
-                            st.session_state.messages.append(message)
-                except RuntimeError as e:
-                    st.error(str(e))
+                with st.spinner("در حال پاسخگویی..."):
+                    try:
+                        # Load embeddings and vector store
+                        embeddings = DeepSeekEmbeddings()
+                        db = FAISS.load_local("faiss_index", embeddings)
+                        docs = db.similarity_search(prompt, k=4)
+                        context = "\n".join([d.page_content for d in docs])
+                        
+                        # Generate response
+                        chat_model = DeepSeekChat()
+                        response = chat_model.generate(context, prompt)
+                        
+                        # Display 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)}")
 
 if __name__ == "__main__":
     main()