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): 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 ) self.embedding_model_name = "intfloat/multilingual-e5-large" self.llm_model_name = "google/flan-t5-large" 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"] ) self.embeddings = HuggingFaceEmbeddings( model_name=self.embedding_model_name, model_kwargs={'device': 'cuda' if torch.cuda.is_available() else 'cpu'} ) self.tokenizer = AutoTokenizer.from_pretrained(self.llm_model_name) self.model = AutoModelForSeq2SeqLM.from_pretrained(self.llm_model_name) self.device = 'cuda' if torch.cuda.is_available() else 'cpu' self.model.to(self.device) self.pipe = pipeline( "text2text-generation", model=self.model, tokenizer=self.tokenizer, max_length=512, device=0 if torch.cuda.is_available() else -1, model_kwargs={"temperature": 0.7} ) self.llm = HuggingFacePipeline(pipeline=self.pipe) def process_documents(self, text: str) -> bool: try: texts = self.text_splitter.split_text(text) self.vectorstore = Chroma.from_texts( texts, self.embeddings, metadatas=[{"source": f"chunk_{i}", "text": t} for i, t in enumerate(texts)], collection_name="enhanced_rag_docs" ) self.retriever = self.vectorstore.as_retriever( search_kwargs={"k": self.k_documents} ) 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} ) return True except Exception as e: print(f"Processing error: {str(e)}") return False def answer_question(self, question: str) -> Tuple[str, str]: try: 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) return answer, sources_text except Exception as e: return f"Error answering: {str(e)}", "" def create_enhanced_interface(): rag_system = EnhancedRAGSystem() def process_and_answer(text: str, question: str) -> str: if not text.strip() or not question.strip(): return "Please provide both text and question." if not rag_system.process_documents(text): return "Error processing the text." answer, sources = rag_system.answer_question(question) if sources: return f"""Answer: {answer} Relevant excerpts consulted: {sources}""" return answer # HTML para o cabeçalho 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("""
""") 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 = [ ["The Earth is the third planet from the Sun. It has one natural satellite called the Moon. It is the only known planet to harbor life.", "What is Earth's natural satellite?"], ["La Tierra es el tercer planeta del Sistema Solar. Tiene un satélite natural llamado Luna. Es el único planeta conocido que alberga vida.", "¿Cuál es el satélite natural de la Tierra?"], ["A Terra é o terceiro planeta do Sistema Solar. Tem um satélite natural chamado Lua. É o único planeta conhecido que abriga vida.", "Qual é o satélite natural da Terra?"], ["The Sun is a medium-sized star at the center of our Solar System. It provides light and heat to all planets.", "What is the Sun?"], ["El Sol es una estrella de tamaño medio en el centro de nuestro Sistema Solar. Proporciona luz y calor a todos los planetas.", "¿Qué es el Sol?"], ["O Sol é uma estrela de tamanho médio no centro do nosso Sistema Solar. Ele fornece luz e calor para todos os planetas.", "O que é o Sol?"] ] 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 ) return interface if __name__ == "__main__": demo = create_enhanced_interface() demo.launch()