Create app.py

app.py

@@ -0,0 +1,362 @@
+import streamlit as st
+import pdfplumber
+import os
+import tempfile
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.vectorstores import FAISS
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from huggingface_hub import InferenceClient
+from langchain.llms.base import LLM
+from typing import Optional, List, Mapping, Any
+import os
+import json
+
+
+class QwenLLM(LLM):
+    client: InferenceClient = None
+
+    def __init__(self):
+        super().__init__()
+        self.client = InferenceClient(
+            provider="hf-inference",
+            api_key="colocar_aqui")
+
+    @property
+    def _llm_type(self) -> str:
+        return "qwen-llm"
+
+    def _call(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[Any] = None,
+        **kwargs: Any,
+    ) -> str:
+        modified_prompt = prompt + "<think>\n"
+        messages = [{"role": "user", "content": modified_prompt}]
+        completion = self.client.chat.completions.create(
+            model="Qwen/QwQ-32B",
+            messages=messages,
+            max_tokens=4096,
+        )
+        return completion.choices[0].message.content
+
+
+# Custom chat history implementation
+class ChatHistory:
+    def __init__(self):
+        self.messages = []
+    
+    def add_user_message(self, message):
+        self.messages.append({"role": "user", "content": message})
+    
+    def add_assistant_message(self, message, sources=None):
+        self.messages.append({
+            "role": "assistant", 
+            "content": message,
+            "sources": sources if sources else []
+        })
+    
+    def get_conversation_history(self, include_sources=False):
+        if include_sources:
+            return self.messages
+        else:
+            # Return messages without sources for sending to LLM
+            return [{"role": m["role"], "content": m["content"]} for m in self.messages]
+    
+    def get_messages_for_display(self):
+        return self.messages
+    
+    def clear(self):
+        self.messages = []
+
+
+# Set page configuration
+st.set_page_config(page_title="RAG Chat with Qwen/QwQ-32B",
+                   page_icon="💬",
+                   layout="wide")
+
+# Initialize session state variables if they don't exist
+if 'vector_store' not in st.session_state:
+    st.session_state.vector_store = None
+if 'document_processed' not in st.session_state:
+    st.session_state.document_processed = False
+if 'file_name' not in st.session_state:
+    st.session_state.file_name = None
+if 'document_text' not in st.session_state:
+    st.session_state.document_text = ""
+if 'chat_history' not in st.session_state:
+    st.session_state.chat_history = ChatHistory()
+
+
+# Function to extract text from document (PDF or TXT)
+def extract_text_from_document(document_file):
+    file_type = document_file.name.split('.')[-1].lower()
+
+    if file_type == 'txt':
+        # For TXT files, simply read the content
+        return document_file.getvalue().decode('utf-8')
+
+    elif file_type == 'pdf':
+        # For PDF files, use pdfplumber
+        with tempfile.NamedTemporaryFile(delete=False,
+                                         suffix='.pdf') as tmp_file:
+            tmp_file.write(document_file.getvalue())
+            tmp_file_path = tmp_file.name
+
+        text = ""
+        try:
+            with pdfplumber.open(tmp_file_path) as pdf:
+                for page in pdf.pages:
+                    page_text = page.extract_text()
+                    if page_text:
+                        text += page_text + "\n\n"
+        except Exception as e:
+            st.error(f"Error extracting text from PDF: {e}")
+        finally:
+            # Clean up the temporary file
+            if os.path.exists(tmp_file_path):
+                os.remove(tmp_file_path)
+
+        return text
+
+    else:
+        st.error(f"Unsupported file type: {file_type}")
+        return ""
+
+
+# Function to create chunks from text
+def create_chunks(text):
+    text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size=500,
+        chunk_overlap=50,
+        length_function=len,
+    )
+    chunks = text_splitter.split_text(text)
+    return chunks
+
+
+# Function to create vector store
+def create_vector_store(chunks):
+    embeddings = HuggingFaceEmbeddings(
+        model_name="sentence-transformers/all-mpnet-base-v2",
+        model_kwargs={'device': 'cpu'})
+    vector_store = FAISS.from_texts(chunks, embeddings)
+    return vector_store
+
+
+# Function to retrieve relevant document chunks
+def retrieve_relevant_chunks(vector_store, query, k=3):
+    if not vector_store:
+        return []
+    docs = vector_store.similarity_search(query, k=k)
+    return docs
+
+
+# Function to generate response using RAG
+def generate_rag_response(query, chat_history, vector_store):
+    # Initialize Qwen LLM
+    llm = QwenLLM()
+    
+    # Retrieve relevant chunks
+    relevant_docs = retrieve_relevant_chunks(vector_store, query, k=3)
+    
+    if not relevant_docs:
+        return "I couldn't find any relevant information in the document to answer your question.", []
+    
+    # Prepare context from relevant documents
+    context = "\n\n".join([doc.page_content for doc in relevant_docs])
+    
+    # Prepare conversation history for context
+    conversation_history = ""
+    for msg in chat_history.get_conversation_history():
+        role = "User" if msg["role"] == "user" else "Assistant"
+        conversation_history += f"{role}: {msg['content']}\n\n"
+    
+    # Create prompt
+    prompt = f"""
+    You are a helpful assistant that provides accurate information based only on the given context and conversation history.
+
+    1. Use only the context below and the conversation history to answer the question.
+    2. If the answer is not in the context, reply with "I don't have enough information to answer this question."
+    3. Be friendly and helpful.
+    4. Maintain continuity with the conversation history.
+
+    Conversation History:
+    {conversation_history}
+
+    Context from document:
+    {context}
+
+    User's question: {query}
+
+    Answer:
+    """
+    
+    # Generate response
+    response = llm(prompt)
+    
+    return response, relevant_docs
+
+
+# Function to handle user message and get AI response
+def process_user_message(user_message):
+    # Add user message to chat history
+    st.session_state.chat_history.add_user_message(user_message)
+    
+    # Generate response
+    with st.spinner("Thinking..."):
+        response, source_docs = generate_rag_response(
+            user_message, 
+            st.session_state.chat_history, 
+            st.session_state.vector_store
+        )
+    
+    # Format sources for storing with the message
+    sources = []
+    for i, doc in enumerate(source_docs):
+        sources.append({"id": i + 1, "content": doc.page_content})
+    
+    # Add assistant response to chat history
+    st.session_state.chat_history.add_assistant_message(response, sources)
+    
+    return response, sources
+
+
+# Main application UI
+st.title("💬 RAG Chat with Qwen/QwQ-32B")
+st.markdown("""
+Upload a PDF or TXT document and chat about its content. This system uses:
+- Document text extraction
+- Text chunking and embedding 
+- Qwen/QwQ-32B for answering questions
+- Memory to maintain conversation context
+""")
+
+# Sidebar for PDF upload and model selection
+with st.sidebar:
+    st.header("Configuration")
+
+    uploaded_file = st.file_uploader("Upload a document", type=['pdf', 'txt'])
+    
+    # Button to clear chat history
+    if st.button("Clear Chat History"):
+        st.session_state.chat_history.clear()
+        st.success("Chat history cleared!")
+
+    st.markdown("**Using Qwen/QwQ-32B model**")
+
+    st.markdown("---")
+    st.markdown("### About")
+    st.markdown("""
+    This is a RAG Chat system that:
+    1. Processes PDF and TXT documents
+    2. Creates a vector database of document content
+    3. Maintains conversation history
+    4. Retrieves relevant information for user queries
+    5. Generates contextual answers using Qwen/QwQ-32B
+    """)
+
+# Process the uploaded document (PDF or TXT)
+if uploaded_file is not None:
+    # Check if we need to process a new file
+    if st.session_state.file_name != uploaded_file.name:
+        st.session_state.file_name = uploaded_file.name
+        st.session_state.document_processed = False
+
+    if not st.session_state.document_processed:
+        file_type = uploaded_file.name.split('.')[-1].lower()
+        with st.spinner(f"Processing {file_type.upper()} file..."):
+            # Extract text from document
+            text = extract_text_from_document(uploaded_file)
+            st.session_state.document_text = text
+
+            # Create chunks from text
+            chunks = create_chunks(text)
+
+            # Create vector store
+            st.session_state.vector_store = create_vector_store(chunks)
+
+            st.session_state.document_processed = True
+            st.success(f"Document processed successfully: {uploaded_file.name}")
+
+            # Display document summary
+            num_chunks = len(chunks)
+            avg_chunk_size = sum(len(chunk) for chunk in chunks) / num_chunks if num_chunks > 0 else 0
+            st.info(f"Document processed into {num_chunks} chunks with average size of {avg_chunk_size:.0f} characters")
+
+# Create two columns for the UI layout
+col1, col2 = st.columns([3, 1])
+
+# Left column for chat interface
+with col1:
+    st.subheader("Chat")
+    
+    # Display chat messages
+    chat_container = st.container()
+    with chat_container:
+        for message in st.session_state.chat_history.get_messages_for_display():
+            with st.chat_message(message["role"]):
+                st.markdown(message["content"])
+                
+                # If the message is from the assistant and has sources, display them
+                if message["role"] == "assistant" and "sources" in message and message["sources"]:
+                    with st.expander("View Sources"):
+                        for source in message["sources"]:
+                            st.markdown(f"**Source {source['id']}**")
+                            st.text(source["content"])
+    
+    # Chat input
+    if st.session_state.document_processed:
+        user_input = st.chat_input("Type your message here...")
+        if user_input:
+            # Display user message
+            with st.chat_message("user"):
+                st.markdown(user_input)
+            
+            # Get and display assistant response
+            response, sources = process_user_message(user_input)
+            with st.chat_message("assistant"):
+                st.markdown(response)
+                
+                # Display sources in an expander
+                if sources:
+                    with st.expander("View Sources"):
+                        for source in sources:
+                            st.markdown(f"**Source {source['id']}**")
+                            st.text(source["content"])
+    else:
+        st.info("Please upload a document to start chatting")
+
+# Right column for document info
+with col2:
+    if st.session_state.document_processed:
+        st.subheader("Document Preview")
+        with st.expander("View Document Text", expanded=False):
+            st.text_area(
+                "Extracted Text",
+                st.session_state.document_text[:5000] +
+                ("..." if len(st.session_state.document_text) > 5000 else ""),
+                height=400)
+    else:
+        st.info("Upload a PDF or TXT document to get started")
+
+# Instructions for users if no document is uploaded
+if not st.session_state.document_processed:
+    st.markdown("""
+    ## Getting Started
+    
+    1. **Upload a PDF or TXT document** using the file uploader in the sidebar
+    2. Wait for the document to be processed
+    3. Start chatting with the AI about the document
+    4. The chat remembers the conversation context
+    5. Clear the chat history using the button in the sidebar
+    
+    The system uses Qwen/QwQ-32B model and maintains conversation memory.
+    """)
+
+# Information about the model
+st.sidebar.info("""
+**Using Hugging Face Inference API**
+This application uses the Qwen/QwQ-32B model via Hugging Face's Inference API for generating responses.
+""")