Update app.py

app.py

@@ -7,86 +7,85 @@ from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_community.document_loaders import TextLoader
 from langchain_huggingface import HuggingFaceEmbeddings
 
-# Environment setup for Torch
+# Fix Torch SDPA issue
 os.environ["TORCH_USE_CUDA_DSA"] = "1"
 
-# Model and tokenizer setup
+# Load Hugging Face model & tokenizer
 MODEL_NAME = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16, device_map="auto")
 
-# Persistent storage paths (Hugging Face Spaces uses /data for persistence)
-CHROMA_DB_PATH = "chroma_db"
-UPLOAD_FOLDER = "uploaded_files"
+# Paths (Hugging Face Spaces uses persistent storage at /data)
+CHROMA_DB_PATH = "/data/chroma_db"
+UPLOAD_FOLDER = "/data/uploaded_files"
 
 # Ensure directories exist
 os.makedirs(UPLOAD_FOLDER, exist_ok=True)
 os.makedirs(CHROMA_DB_PATH, exist_ok=True)
 
-# Initialize ChromaDB with embeddings
+# Load embedding function
 embedding_function = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+
+# Load ChromaDB (persistent across sessions)
 vector_db = Chroma(persist_directory=CHROMA_DB_PATH, embedding_function=embedding_function)
 
-# Prompt template
+# Prompt Template
 PROMPT_TEMPLATE = """
-You are an expert research assistant. Use the provided context to answer the query concisely (max 3 sentences).
-Query: {user_query}
-Context: {document_context}
+You are an expert research assistant. Use the provided context to answer the query. 
+If unsure, state that you don't know. Be concise and factual (max 3 sentences).
+Query: {user_query} 
+Context: {document_context} 
 Answer:
 """
 
-# Process and store uploaded text file
+# Function to process and store uploaded file
 def process_and_store_file(file_path):
-    try:
-        loader = TextLoader(file_path)
-        raw_docs = loader.load()
-        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-        document_chunks = text_splitter.split_documents(raw_docs)
-        vector_db.add_documents(document_chunks)
-        vector_db.persist()
-        return "File processed and stored successfully."
-    except Exception as e:
-        return f"Error processing file: {str(e)}"
-
-# Search ChromaDB for relevant documents
+    loader = TextLoader(file_path)
+    raw_docs = loader.load()
+    
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    document_chunks = text_splitter.split_documents(raw_docs)
+    
+    vector_db.add_documents(document_chunks)
+
+# Function to search documents in ChromaDB
 def find_related_documents(query):
     return vector_db.similarity_search(query, k=3)
 
-# Generate answer using the model
+# Function to generate answers using Hugging Face Model
 def generate_answer(question, uploaded_file=None):
-    # Handle file upload if provided
-    if uploaded_file:
+    # Process file if uploaded
+    if uploaded_file is not None:
         file_path = os.path.join(UPLOAD_FOLDER, uploaded_file.name)
         with open(file_path, "wb") as f:
             f.write(uploaded_file.read())
         process_and_store_file(file_path)
 
-    # Retrieve relevant context
+    # Search for relevant documents
     relevant_docs = find_related_documents(question)
-    context = "\n".join([doc.page_content for doc in relevant_docs]) if relevant_docs else "No relevant context found."
+    context = "\n\n".join([doc.page_content for doc in relevant_docs]) if relevant_docs else "No relevant documents found."
 
-    # Format prompt
+    # Prepare prompt
     prompt = PROMPT_TEMPLATE.format(user_query=question, document_context=context)
+
+    # Tokenize & generate response
     inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
-    outputs = model.generate(**inputs, max_new_tokens=150, temperature=0.7)
+    outputs = model.generate(**inputs, max_length=200, num_return_sequences=1, temperature=0.7)
     answer = tokenizer.decode(outputs[0], skip_special_tokens=True).split("Answer:")[-1].strip()
 
     return answer
 
-# Gradio interface
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("### 📘 DocuMind AI - Powered by DeepSeek Qwen")
-    gr.Markdown("Upload a text file to build a knowledge base, then ask questions based on it.")
+# Gradio UI
+with gr.Blocks(title="DocuMind AI - DeepSeek Qwen Chatbot") as demo:
+    gr.Markdown("### 📘 DocuMind AI - Ask Questions Based on Uploaded Documents")
     
     with gr.Row():
-        with gr.Column(scale=2):
-            question_input = gr.Textbox(label="Your Question", placeholder="Type your question here...")
-        with gr.Column(scale=1):
-            file_input = gr.File(label="Upload Text File (Optional)", file_types=[".txt"])
-    
-    submit_btn = gr.Button("Submit", variant="primary")
-    output_text = gr.Textbox(label="Answer", lines=5, placeholder="Answer will appear here...")
+        question_input = gr.Textbox(label="Ask a Question", placeholder="Type your question here...")
+        file_input = gr.File(label="Upload a Text File (Optional)", file_types=[".txt"])
+
+    submit_btn = gr.Button("Get Answer")
+    output_text = gr.Textbox(label="Answer", interactive=False)
 
-    submit_btn.click(fn=generate_answer, inputs=[question_input, file_input], outputs=output_text)
+    submit_btn.click(generate_answer, inputs=[question_input, file_input], outputs=output_text)
 
 demo.launch()