Update app.py

app.py

@@ -1,4 +1,4 @@
-import streamlit as st
+import gradio as gr
 import os
 import re
 import torch
@@ -12,138 +12,64 @@ from langchain_community.vectorstores import FAISS
 from langchain.schema import Document
 from langchain.embeddings import HuggingFaceEmbeddings
 
-# Set page configuration
-st.set_page_config(
-    page_title="Vision 2030 Virtual Assistant",
-    page_icon="🇸🇦",
-    layout="wide"
-)
-
-# App title and description
-st.title("Vision 2030 Virtual Assistant")
-st.markdown("Ask questions about Saudi Vision 2030 goals, projects, and progress in Arabic or English.")
-
-# Function definitions
-@st.cache_resource
-def load_model_and_tokenizer():
-    """Load the ALLaM-7B model and tokenizer with error handling"""
-    model_name = "ALLaM-AI/ALLaM-7B-Instruct-preview"
-    st.info(f"Loading model: {model_name} (this may take a few minutes)")
-    
-    try:
-        # First attempt with AutoTokenizer
-        tokenizer = AutoTokenizer.from_pretrained(
-            model_name,
-            trust_remote_code=True,
-            use_fast=False
-        )
+# Create the Vision 2030 Assistant class
+class Vision2030Assistant:
+    def __init__(self, model, tokenizer, vector_store):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.vector_store = vector_store
+        self.conversation_history = []
         
-        # Load model with appropriate settings for ALLaM
-        model = AutoModelForCausalLM.from_pretrained(
-            model_name,
-            torch_dtype=torch.bfloat16,
-            trust_remote_code=True,
-            device_map="auto",
-        )
+    def answer(self, user_query):
+        # Detect language
+        language = detect_language(user_query)
         
-        st.success("Model loaded successfully!")
+        # Add user query to conversation history
+        self.conversation_history.append({"role": "user", "content": user_query})
         
-    except Exception as e:
-        st.error(f"First loading attempt failed: {e}")
-        st.info("Trying alternative loading approach...")
+        # Get the full conversation context
+        conversation_context = "\n".join([
+            f"{'User' if msg['role'] == 'user' else 'Assistant'}: {msg['content']}"
+            for msg in self.conversation_history[-6:]  # Keep last 3 turns (6 messages)
+        ])
         
-        # Try with specific tokenizer class if the first attempt fails
-        from transformers import LlamaTokenizer
+        # Enhance query with conversation context for better retrieval
+        enhanced_query = f"{conversation_context}\n{user_query}"
         
-        tokenizer = LlamaTokenizer.from_pretrained(model_name)
-        model = AutoModelForCausalLM.from_pretrained(
-            model_name,
-            torch_dtype=torch.float16,
-            trust_remote_code=True,
-            device_map="auto",
-        )
+        # Retrieve relevant contexts
+        contexts = retrieve_context(enhanced_query, self.vector_store, top_k=5)
+        
+        # Generate response
+        response = generate_response(user_query, contexts, self.model, self.tokenizer, language)
+        
+        # Add response to conversation history
+        self.conversation_history.append({"role": "assistant", "content": response})
+        
+        # Also return sources for transparency
+        sources = [ctx.get("source", "Unknown") for ctx in contexts]
+        unique_sources = list(set(sources))
         
-        st.success("Model loaded successfully with LlamaTokenizer!")
+        # Format the response with sources
+        if unique_sources:
+            source_text = "\n\nSources: " + ", ".join([os.path.basename(src) for src in unique_sources])
+            response_with_sources = response + source_text
+        else:
+            response_with_sources = response
+        
+        return response_with_sources
     
-    return model, tokenizer
+    def reset_conversation(self):
+        """Reset the conversation history"""
+        self.conversation_history = []
+        return "Conversation has been reset."
 
+# Helper functions
 def detect_language(text):
     """Detect if text is primarily Arabic or English"""
     arabic_chars = re.findall(r'[\u0600-\u06FF]', text)
     is_arabic = len(arabic_chars) > len(text) * 0.5
     return "arabic" if is_arabic else "english"
 
-def process_pdfs():
-    """Process uploaded PDF documents"""
-    documents = []
-    
-    if 'uploaded_pdfs' in st.session_state and st.session_state.uploaded_pdfs:
-        for pdf_file in st.session_state.uploaded_pdfs:
-            try:
-                # Save the uploaded file temporarily
-                pdf_path = f"temp_{pdf_file.name}"
-                with open(pdf_path, "wb") as f:
-                    f.write(pdf_file.getbuffer())
-                
-                # Extract text
-                text = ""
-                with open(pdf_path, 'rb') as file:
-                    reader = PyPDF2.PdfReader(file)
-                    for page in reader.pages:
-                        text += page.extract_text() + "\n\n"
-                
-                # Remove temporary file
-                os.remove(pdf_path)
-                
-                if text.strip():  # If we got some text
-                    doc = Document(
-                        page_content=text,
-                        metadata={"source": pdf_file.name, "filename": pdf_file.name}
-                    )
-                    documents.append(doc)
-                    st.info(f"Successfully processed: {pdf_file.name}")
-                else:
-                    st.warning(f"No text extracted from {pdf_file.name}")
-            except Exception as e:
-                st.error(f"Error processing {pdf_file.name}: {e}")
-    
-    st.success(f"Processed {len(documents)} PDF documents")
-    return documents
-
-def create_vector_store(documents):
-    """Split documents into chunks and create a FAISS vector store"""
-    # Text splitter for breaking documents into chunks
-    text_splitter = RecursiveCharacterTextSplitter(
-        chunk_size=500,
-        chunk_overlap=50,
-        separators=["\n\n", "\n", ".", "!", "?", ",", " ", ""]
-    )
-    
-    # Split documents into chunks
-    chunks = []
-    for doc in documents:
-        doc_chunks = text_splitter.split_text(doc.page_content)
-        # Preserve metadata for each chunk
-        chunks.extend([
-            Document(page_content=chunk, metadata=doc.metadata)
-            for chunk in doc_chunks
-        ])
-    
-    st.info(f"Created {len(chunks)} chunks from {len(documents)} documents")
-    
-    # Create a proper embedding function for LangChain
-    embedding_function = HuggingFaceEmbeddings(
-        model_name="sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
-    )
-    
-    # Create FAISS index
-    vector_store = FAISS.from_documents(
-        chunks,
-        embedding_function
-    )
-    
-    return vector_store
-
 def retrieve_context(query, vector_store, top_k=5):
     """Retrieve most relevant document chunks for a given query"""
     # Search the vector store using similarity search
@@ -160,10 +86,11 @@ def retrieve_context(query, vector_store, top_k=5):
     
     return contexts
 
-def generate_response(query, contexts, model, tokenizer):
+def generate_response(query, contexts, model, tokenizer, language="auto"):
     """Generate a response using retrieved contexts with ALLaM-specific formatting"""
-    # Auto-detect language
-    language = detect_language(query)
+    # Auto-detect language if not specified
+    if language == "auto":
+        language = detect_language(query)
     
     # Format the prompt based on language
     if language == "arabic":
@@ -189,113 +116,262 @@ Context:
 Question: {query} [/INST]</s>"""
     
     try:
-        with st.spinner("Generating response..."):
-            # Generate response with appropriate parameters for ALLaM
-            inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-            
-            # Generate with appropriate parameters
-            outputs = model.generate(
-                inputs.input_ids,
-                attention_mask=inputs.attention_mask,
-                max_new_tokens=512,
-                temperature=0.7,
-                top_p=0.9,
-                do_sample=True,
-                repetition_penalty=1.1
-            )
-            
-            # Decode the response
-            full_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
-            
-            # Extract just the answer part (after the instruction)
-            response = full_output.split("[/INST]")[-1].strip()
+        # Generate response with appropriate parameters for ALLaM
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        
+        # Generate with appropriate parameters
+        outputs = model.generate(
+            inputs.input_ids,
+            attention_mask=inputs.attention_mask,
+            max_new_tokens=512,
+            temperature=0.7,
+            top_p=0.9,
+            do_sample=True,
+            repetition_penalty=1.1
+        )
+        
+        # Decode the response
+        full_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # Extract just the answer part (after the instruction)
+        response = full_output.split("[/INST]")[-1].strip()
+        
+        # If response is empty for some reason, return the full output
+        if not response:
+            response = full_output
             
-            # If response is empty for some reason, return the full output
-            if not response:
-                response = full_output
-                
-            return response, [ctx.get("source", "Unknown") for ctx in contexts]
+        return response
         
     except Exception as e:
-        st.error(f"Error during generation: {e}")
+        print(f"Error during generation: {e}")
         # Fallback response
-        return "I apologize, but I encountered an error while generating a response.", []
-
-# Initialize the app state
-if 'conversation_history' not in st.session_state:
-    st.session_state.conversation_history = []
-
-if 'vector_store' not in st.session_state:
-    st.session_state.vector_store = None
-
-if 'uploaded_pdfs' not in st.session_state:
-    st.session_state.uploaded_pdfs = None
+        return "I apologize, but I encountered an error while generating a response."
 
-# PDF upload section
-st.header("1. Upload Vision 2030 Documents")
-uploaded_files = st.file_uploader("Upload PDF documents about Vision 2030", 
-                                 type=["pdf"], 
-                                 accept_multiple_files=True,
-                                 help="Upload one or more PDF documents containing information about Vision 2030")
-
-if uploaded_files:
-    st.session_state.uploaded_pdfs = uploaded_files
-    if st.button("Process PDFs"):
-        documents = process_pdfs()
-        if documents:
-            with st.spinner("Creating vector database..."):
-                st.session_state.vector_store = create_vector_store(documents)
-            st.success("Vector database created successfully!")
+def process_pdf_files(pdf_files):
+    """Process PDF files and create documents"""
+    documents = []
+    
+    for pdf_file in pdf_files:
+        try:
+            # Save the uploaded file temporarily
+            temp_path = f"temp_{pdf_file.name}"
+            with open(temp_path, "wb") as f:
+                f.write(pdf_file.read())
+            
+            # Extract text
+            text = ""
+            with open(temp_path, 'rb') as file:
+                reader = PyPDF2.PdfReader(file)
+                for page in reader.pages:
+                    page_text = page.extract_text()
+                    if page_text:
+                        text += page_text + "\n\n"
+            
+            # Clean up
+            os.remove(temp_path)
+            
+            if text.strip():  # If we got some text
+                doc = Document(
+                    page_content=text,
+                    metadata={"source": pdf_file.name, "filename": pdf_file.name}
+                )
+                documents.append(doc)
+                print(f"Successfully processed: {pdf_file.name}")
+            else:
+                print(f"Warning: No text extracted from {pdf_file.name}")
+        except Exception as e:
+            print(f"Error processing {pdf_file.name}: {e}")
+    
+    print(f"Processed {len(documents)} PDF documents")
+    return documents
 
-# Load the model (cached)
-model, tokenizer = load_model_and_tokenizer()
+def create_vector_store(documents):
+    """Create a vector store from documents"""
+    # Text splitter for breaking documents into chunks
+    text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size=500,
+        chunk_overlap=50,
+        separators=["\n\n", "\n", ".", "!", "?", ",", " ", ""]
+    )
+    
+    # Split documents into chunks
+    chunks = []
+    for doc in documents:
+        doc_chunks = text_splitter.split_text(doc.page_content)
+        # Preserve metadata for each chunk
+        chunks.extend([
+            Document(page_content=chunk, metadata=doc.metadata)
+            for chunk in doc_chunks
+        ])
+    
+    print(f"Created {len(chunks)} chunks from {len(documents)} documents")
+    
+    # Create embedding function
+    embedding_function = HuggingFaceEmbeddings(
+        model_name="sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
+    )
+    
+    # Create FAISS index
+    vector_store = FAISS.from_documents(chunks, embedding_function)
+    return vector_store
 
-# Chat interface
-st.header("2. Chat with the Vision 2030 Assistant")
+# Variables to store state
+model = None
+tokenizer = None
+assistant = None
 
-# Display conversation history
-for message in st.session_state.conversation_history:
-    if message["role"] == "user":
-        st.markdown(f"**You:** {message['content']}")
-    else:
-        st.markdown(f"**Assistant:** {message['content']}")
-        if 'sources' in message and message['sources']:
-            st.markdown(f"*Sources: {', '.join([os.path.basename(src) for src in message['sources']])}*")
-    st.divider()
+# Load the model and tokenizer
+def load_model_and_tokenizer():
+    global model, tokenizer
+    
+    if model is not None and tokenizer is not None:
+        return "Model already loaded"
+    
+    model_name = "ALLaM-AI/ALLaM-7B-Instruct-preview"
+    print(f"Loading model: {model_name}")
 
-# Input for new question
-user_input = st.text_input("Ask a question about Vision 2030 (in Arabic or English):", key="user_query")
+    try:
+        # First attempt with AutoTokenizer
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_name,
+            trust_remote_code=True,
+            use_fast=False
+        )
+        
+        # Load model with appropriate settings for ALLaM
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.bfloat16,  # Use bfloat16 for better compatibility
+            trust_remote_code=True,
+            device_map="auto",
+        )
+        
+        return "Model loaded successfully with AutoTokenizer!"
+        
+    except Exception as e:
+        error_msg = f"First loading attempt failed: {e}"
+        print(error_msg)
+        
+        try:
+            # Try with specific tokenizer class if the first attempt fails
+            from transformers import LlamaTokenizer
+            
+            tokenizer = LlamaTokenizer.from_pretrained(model_name)
+            model = AutoModelForCausalLM.from_pretrained(
+                model_name,
+                torch_dtype=torch.float16,
+                trust_remote_code=True,
+                device_map="auto",
+            )
+            
+            return "Model loaded successfully with LlamaTokenizer!"
+        except Exception as e2:
+            return f"Both loading attempts failed. Error 1: {e}. Error 2: {e2}"
 
-# Examples
-st.markdown("**Example questions:**")
-examples_col1, examples_col2 = st.columns(2)
-with examples_col1:
-    st.markdown("- What is Saudi Vision 2030?")
-    st.markdown("- What are the economic goals of Vision 2030?")
-    st.markdown("- How does Vision 2030 support women's empowerment?")
-with examples_col2:
-    st.markdown("- ما هي رؤية السعودية 2030؟")
-    st.markdown("- ما هي الأهداف الاقتصادية لرؤية 2030؟")
-    st.markdown("- كيف تدعم رؤية 2030 تمكين المرأة السعودية؟")
+# Gradio Interface Functions
+def process_pdfs(pdf_files):
+    if not pdf_files:
+        return "No files uploaded. Please upload PDF documents about Vision 2030."
+    
+    documents = process_pdf_files(pdf_files)
+    
+    if not documents:
+        return "Failed to extract text from the uploaded PDFs."
+    
+    global assistant, model, tokenizer
+    
+    # Ensure model is loaded
+    if model is None or tokenizer is None:
+        load_status = load_model_and_tokenizer()
+        if "successfully" not in load_status.lower():
+            return f"Model loading failed: {load_status}"
+    
+    # Create vector store
+    vector_store = create_vector_store(documents)
+    
+    # Initialize assistant
+    assistant = Vision2030Assistant(model, tokenizer, vector_store)
+    
+    return f"Successfully processed {len(documents)} documents. The assistant is ready to use!"
 
-# Process the user input
-if user_input and st.session_state.vector_store:
-    # Add user message to history
-    st.session_state.conversation_history.append({"role": "user", "content": user_input})
+def answer_query(message, history):
+    global assistant
     
-    # Get response
-    response, sources = generate_response(user_input, retrieve_context(user_input, st.session_state.vector_store), model, tokenizer)
+    if assistant is None:
+        return "Please upload and process Vision 2030 PDF documents first."
     
-    # Add assistant message to history
-    st.session_state.conversation_history.append({"role": "assistant", "content": response, "sources": sources})
+    response = assistant.answer(message)
+    return response
+
+def reset_chat():
+    global assistant
+    
+    if assistant is None:
+        return "No active conversation to reset."
     
-    # Rerun to update the UI
-    st.experimental_rerun()
+    reset_message = assistant.reset_conversation()
+    return reset_message
 
-elif user_input and not st.session_state.vector_store:
-    st.warning("Please upload and process Vision 2030 PDF documents first")
+# Create Gradio interface
+with gr.Blocks(title="Vision 2030 Virtual Assistant") as demo:
+    gr.Markdown("# Vision 2030 Virtual Assistant")
+    gr.Markdown("Ask questions about Saudi Vision 2030 goals, projects, and progress in Arabic or English.")
+    
+    with gr.Tab("Setup"):
+        gr.Markdown("## Step 1: Load the Model")
+        load_btn = gr.Button("Load ALLaM-7B Model", variant="primary")
+        load_output = gr.Textbox(label="Load Status")
+        load_btn.click(load_model_and_tokenizer, inputs=[], outputs=load_output)
+        
+        gr.Markdown("## Step 2: Upload Vision 2030 Documents")
+        pdf_files = gr.File(file_types=[".pdf"], file_count="multiple", label="Upload PDF Documents")
+        process_btn = gr.Button("Process Documents", variant="primary")
+        process_output = gr.Textbox(label="Processing Status")
+        process_btn.click(process_pdfs, inputs=[pdf_files], outputs=process_output)
+    
+    with gr.Tab("Chat"):
+        chatbot = gr.Chatbot(label="Conversation")
+        message = gr.Textbox(
+            label="Ask a question about Vision 2030 (in Arabic or English)",
+            placeholder="What are the main goals of Vision 2030?",
+            lines=2
+        )
+        submit_btn = gr.Button("Submit", variant="primary")
+        reset_btn = gr.Button("Reset Conversation")
+        
+        gr.Markdown("### Example Questions")
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("**English Questions:**")
+                en_examples = gr.Examples(
+                    examples=[
+                        "What is Saudi Vision 2030?",
+                        "What are the economic goals of Vision 2030?",
+                        "How does Vision 2030 support women's empowerment?",
+                        "What environmental initiatives are part of Vision 2030?",
+                        "What is the role of the Public Investment Fund in Vision 2030?"
+                    ],
+                    inputs=message
+                )
+            
+            with gr.Column():
+                gr.Markdown("**Arabic Questions:**")
+                ar_examples = gr.Examples(
+                    examples=[
+                        "ما هي رؤية السعودية 2030؟",
+                        "ما هي الأهداف الاقتصادية لرؤية 2030؟",
+                        "كيف تدعم رؤية 2030 تمكين المرأة السعودية؟",
+                        "ما هي مبادرات رؤية 2030 للحفاظ على البيئة؟",
+                        "ما هي استراتيجية صندوق الاستثمارات العامة في رؤية 2030؟"
+                    ],
+                    inputs=message
+                )
+        
+        reset_output = gr.Textbox(label="Reset Status", visible=False)
+        submit_btn.click(answer_query, inputs=[message, chatbot], outputs=[chatbot])
+        message.submit(answer_query, inputs=[message, chatbot], outputs=[chatbot])
+        reset_btn.click(reset_chat, inputs=[], outputs=[reset_output])
+        reset_btn.click(lambda: None, inputs=[], outputs=[chatbot], postprocess=False)
 
-# Reset conversation button
-if st.button("Reset Conversation") and len(st.session_state.conversation_history) > 0:
-    st.session_state.conversation_history = []
-    st.experimental_rerun()
+# Launch the app
+demo.launch()