Update app.py

app.py

@@ -8,97 +8,83 @@ from langchain.chains import RetrievalQA
 from langchain_community.llms import HuggingFacePipeline
 from transformers import pipeline, AutoTokenizer
 
-# Configuration
-DOCS_FOLDER = "study_materials"
-CHUNK_SIZE = 1000
-CHUNK_OVERLAP = 150
-MODEL_NAME = "google/flan-t5-base"
-
-def get_documents():
-    """Load and process documents without external dependencies"""
+def load_documents(file_path="study_materials"):
     documents = []
-    for file in os.listdir(DOCS_FOLDER):
-        path = os.path.join(DOCS_FOLDER, file)
-        if file.endswith(".pdf"):
+    for filename in os.listdir(file_path):
+        path = os.path.join(file_path, filename)
+        if filename.endswith(".pdf"):
             loader = PyMuPDFLoader(path)
             documents.extend(loader.load())
-        elif file.endswith(".txt"):
+        elif filename.endswith(".txt"):
             loader = TextLoader(path)
             documents.extend(loader.load())
     return documents
 
-def initialize_system():
+def create_qa_system():
     try:
-        # 1. Document Processing
-        docs = get_documents()
-        if not docs:
-            raise RuntimeError(f"⚠️ No documents found in {DOCS_FOLDER} folder")
+        # Load documents
+        documents = load_documents()
+        if not documents:
+            raise ValueError("📚 No study materials found")
         
-        # 2. Text Chunking
-        splitter = CharacterTextSplitter(
-            chunk_size=CHUNK_SIZE,
-            chunk_overlap=CHUNK_OVERLAP,
+        # Text splitting
+        text_splitter = CharacterTextSplitter(
+            chunk_size=800,
+            chunk_overlap=100,
             separator="\n\n"
         )
-        chunks = splitter.split_documents(docs)
+        texts = text_splitter.split_documents(documents)
         
-        # 3. Local Embeddings
+        # Embeddings
         embeddings = HuggingFaceEmbeddings(
             model_name="sentence-transformers/all-MiniLM-L6-v2"
         )
         
-        # 4. Vector Store
-        vector_db = FAISS.from_documents(chunks, embeddings)
+        # Vector store
+        db = FAISS.from_documents(texts, embeddings)
         
-        # 5. Local LLM Setup
-        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-        text_gen = pipeline(
-            task="text2text-generation",
-            model=MODEL_NAME,
+        # LLM setup with proper LangChain wrapper
+        tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
+        pipe = pipeline(
+            "text2text-generation",
+            model="google/flan-t5-base",
             tokenizer=tokenizer,
-            max_length=500,
-            temperature=0.4,
-            device=-1  # Force CPU
+            max_length=300,
+            temperature=0.3,
+            device=-1
         )
         
-        # 6. LangChain Integration
-        llm = HuggingFacePipeline(pipeline=text_gen)
+        # Wrap pipeline in LangChain component
+        llm = HuggingFacePipeline(pipeline=pipe)
         
-        return RetrievalQA.from_chain_type(
+        # Create QA chain
+        return RetrievalQA.from_llm(
             llm=llm,
-            chain_type="stuff",
-            retriever=vector_db.as_retriever(search_kwargs={"k": 3}),
+            retriever=db.as_retriever(search_kwargs={"k": 2}),
             return_source_documents=True
         )
-    except Exception as error:
-        raise RuntimeError(f"Initialization failed: {str(error)}")
+    except Exception as e:
+        raise gr.Error(f"Error: {str(e)}")
 
-# Initialize QA system
+# Initialize system
 try:
-    qa_system = initialize_system()
-except Exception as error:
-    print(f"Fatal Error: {str(error)}")
+    qa = create_qa_system()
+except Exception as e:
+    print(f"Startup failed: {str(e)}")
     raise
 
-def handle_query(query, history):
-    """Process user queries with enhanced error handling"""
+def ask_question(question, history):
     try:
-        result = qa_system.invoke({"query": query})
-        response = result["result"]
-        sources = {doc.metadata['source'] for doc in result['source_documents']}
-        return f"{response}\n\nSources: {', '.join(sources)}"
-    except Exception as error:
-        print(f"Query Error: {str(error)}")
-        return "Error processing request. Please check document formatting."
+        result = qa.invoke({"query": question})
+        answer = result["result"]
+        sources = list({doc.metadata['source'] for doc in result['source_documents']})
+        return f"{answer}\n\n📚 Sources: {', '.join(sources)}"
+    except Exception as e:
+        return f"Error: {str(e)[:150]}"
 
-# Create interface
 gr.ChatInterface(
-    fn=handle_query,
-    title="Local Document AI",
-    description="Upload PDF/TXT files to 'study_materials' folder and ask questions",
-    examples=[
-        "Summarize the main points from chapter 3",
-        "Explain the key concepts in section 2.1",
-        "What are the advantages discussed on page 4?"
-    ]
+    ask_question,
+    title="Study Assistant",
+    description="Upload PDF/TXT files in 'study_materials' folder and ask questions!",
+    theme="soft"
 ).launch()