app.py

import os
from typing import List, Tuple, Dict
from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer
from sentence_transformers import SentenceTransformer
from langchain_community.vectorstores import Chroma
from langchain.chains import RetrievalQA
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.llms import HuggingFacePipeline
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.prompts import PromptTemplate
import gradio as gr
import torch

class EnhancedRAGSystem:
    def __init__(self):
        try:
            print("Initializing RAG System...")
            self.chunk_size = 500
            self.chunk_overlap = 50
            self.k_documents = 4
            
            self.text_splitter = RecursiveCharacterTextSplitter(
                chunk_size=self.chunk_size,
                chunk_overlap=self.chunk_overlap,
                length_function=len
            )
            
            print("Loading embedding model...")
            self.embeddings = HuggingFaceEmbeddings(
                model_name="sentence-transformers/all-MiniLM-L6-v2",
                model_kwargs={'device': 'cpu'}
            )
            
            print("Loading language model...")
            self.llm_model_name = "google/flan-t5-small"
            self.tokenizer = AutoTokenizer.from_pretrained(self.llm_model_name)
            self.model = AutoModelForSeq2SeqLM.from_pretrained(self.llm_model_name)
            
            self.prompt_template = PromptTemplate(
                template="""Use the context below to answer the question. 
                If the answer is not in the context, say "I don't have enough information in the context to answer this question."
                
                Context: {context}
                Question: {question}
                
                Detailed answer:""",
                input_variables=["context", "question"]
            )
            
            print("Setting up pipeline...")
            self.pipe = pipeline(
                "text2text-generation",
                model=self.model,
                tokenizer=self.tokenizer,
                max_length=512,
                device=-1,
                model_kwargs={"temperature": 0.7}
            )
            
            self.llm = HuggingFacePipeline(pipeline=self.pipe)
            print("RAG System initialized successfully!")
            
        except Exception as e:
            print(f"Error during initialization: {str(e)}")
            raise

    def process_documents(self, text: str) -> bool:
        try:
            print("Processing documents...")
            if not text or len(text.strip()) < 10:
                print("Text is too short or empty")
                return False
                
            print("Splitting text...")
            texts = self.text_splitter.split_text(text)
            
            print("Creating vectorstore...")
            self.vectorstore = Chroma.from_texts(
                texts,
                self.embeddings,
                metadatas=[{"source": f"chunk_{i}", "text": t} for i, t in enumerate(texts)]
            )
            
            print("Setting up retriever...")
            self.retriever = self.vectorstore.as_retriever(
                search_kwargs={"k": self.k_documents}
            )
            
            print("Creating QA chain...")
            self.qa_chain = RetrievalQA.from_chain_type(
                llm=self.llm,
                chain_type="stuff",
                retriever=self.retriever,
                return_source_documents=True,
                chain_type_kwargs={"prompt": self.prompt_template}
            )
            
            print("Documents processed successfully!")
            return True
            
        except Exception as e:
            print(f"Error processing documents: {str(e)}")
            return False

    def answer_question(self, question: str) -> Tuple[str, str]:
        try:
            print(f"Answering question: {question}")
            if not hasattr(self, 'qa_chain'):
                return "Please process some documents first.", ""
                
            response = self.qa_chain({"query": question})
            answer = response["result"]
            
            sources = []
            for i, doc in enumerate(response["source_documents"], 1):
                text_preview = doc.page_content[:100] + "..."
                sources.append(f"Excerpt {i}: {text_preview}")
            
            sources_text = "\n".join(sources)
            print("Answer generated successfully!")
            return answer, sources_text
            
        except Exception as e:
            print(f"Error generating answer: {str(e)}")
            return f"Error generating answer: {str(e)}", ""

def create_enhanced_interface():
    try:
        print("Creating interface...")
        rag_system = EnhancedRAGSystem()
        
        def process_and_answer(text: str, question: str) -> str:
            print("Processing new request...")
            if not text.strip() or not question.strip():
                return "Please provide both text and question."
            
            success = rag_system.process_documents(text)
            if not success:
                return "Error processing the text. Please check if the text is valid and try again."
            
            answer, sources = rag_system.answer_question(question)
            
            if sources:
                return f"""Answer: {answer}

Relevant excerpts consulted:
{sources}"""
            return answer

        custom_css = """
            .custom-description {
                margin-bottom: 20px;
                text-align: center;
            }
            .custom-description a {
                text-decoration: none;
                color: #007bff;
                margin: 0 5px;
            }
            .custom-description a:hover {
                text-decoration: underline;
            }
        """

        with gr.Blocks(css=custom_css) as interface:
            gr.HTML("""
                <div class="custom-description">
                    <h1>Advanced RAG with Multilingual Support</h1>
                    <p>Ramon Mayor Martins: 
                        <a href="https://rmayormartins.github.io/" target="_blank">Website</a> | 
                        <a href="https://huggingface.co/rmayormartins" target="_blank">Spaces</a> |
                        <a href="https://github.com/rmayormartins" target="_blank">GitHub</a>
                    </p>
                    <p>This system uses Retrieval-Augmented Generation (RAG) to answer questions about your texts in multiple languages.
                    Simply paste your text and ask questions in any language!</p>
                </div>
            """)
            
            with gr.Row():
                with gr.Column():
                    text_input = gr.Textbox(
                        label="Base Text",
                        placeholder="Paste here the text that will serve as knowledge base...",
                        lines=10
                    )
                    question_input = gr.Textbox(
                        label="Your Question",
                        placeholder="What would you like to know about the text?"
                    )
                    submit_btn = gr.Button("Submit")
                
                with gr.Column():
                    output = gr.Textbox(label="Answer")

            examples = [
                # English example
                ["The solar system consists of the Sun and the celestial bodies that orbit it. These include eight planets (Mercury, Venus, Earth, Mars, Jupiter, Saturn, Uranus, and Neptune), their moons, asteroids, comets, and other objects.",
                 "How many planets are in the solar system?"],
                
                # Spanish example
                ["El sistema solar está formado por el Sol y los cuerpos celestes que orbitan a su alrededor. Estos incluyen ocho planetas (Mercurio, Venus, Tierra, Marte, Júpiter, Saturno, Urano y Neptuno), sus lunas, asteroides, cometas y otros objetos.",
                 "¿Cuántos planetas hay en el sistema solar?"],
                
                # Portuguese example
                ["O sistema solar é composto pelo Sol e pelos corpos celestes que orbitam ao seu redor. Isso inclui oito planetas (Mercúrio, Vênus, Terra, Marte, Júpiter, Saturno, Urano e Netuno), suas luas, asteroides, cometas e outros objetos.",
                 "Quantos planetas existem no sistema solar?"]
            ]

            gr.Examples(
                examples=examples,
                inputs=[text_input, question_input],
                outputs=output,
                fn=process_and_answer,
                cache_examples=True
            )

            submit_btn.click(
                fn=process_and_answer,
                inputs=[text_input, question_input],
                outputs=output
            )

        print("Interface created successfully!")
        return interface
        
    except Exception as e:
        print(f"Error creating interface: {str(e)}")
        raise

if __name__ == "__main__":
    print("Starting application...")
    try:
        demo = create_enhanced_interface()
        demo.launch()
    except Exception as e:
        print(f"Application failed to start: {str(e)}")