Update app.py

app.py

@@ -1,35 +1,34 @@
 import streamlit as st
-import openai
-from langchain.docstore.document import Document
-from langchain.text_splitter import CharacterTextSplitter
-from langchain.chains.summarize import load_summarize_chain
-from langchain.llms import OpenAI
+from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration, pipeline
 
-def generate_response(txt, openai_api_key):
+# Function to generate response using RAG (Retrieval-Augmented Generation)
+def generate_response_with_rag(txt):
     try:
-        # Set up OpenAI API key
-        openai.api_key = openai_api_key
+        # Initialize the RAG model and tokenizer
+        tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-nq")
+        retriever = RagRetriever.from_pretrained("facebook/rag-token-nq", index_name="exact", use_dummy_dataset=True)
+        model = RagSequenceForGeneration.from_pretrained("facebook/rag-token-nq")
 
-        # Split text into chunks
-        text_splitter = CharacterTextSplitter()
-        texts = text_splitter.split_text(txt)
+        # Tokenize the input text
+        inputs = tokenizer(txt, return_tensors="pt")
 
-        # Create documents from the text chunks
-        docs = [Document(page_content=t) for t in texts]
+        # Retrieve relevant documents using the retriever
+        retrieved_docs = retriever.retrieve(inputs["input_ids"])
 
-        # Instantiate the OpenAI LLM model
-        llm = OpenAI(temperature=0, openai_api_key=openai_api_key)
+        # Generate the output using RAG
+        generated = model.generate(input_ids=inputs["input_ids"], context_input_ids=retrieved_docs['context_input_ids'])
 
-        # Create the summarization chain and summarize the documents
-        chain = load_summarize_chain(llm=llm, chain_type='map_reduce')
-        return chain.run(docs)
+        # Decode the generated text
+        summary = tokenizer.decode(generated[0], skip_special_tokens=True)
+
+        return summary
     except Exception as e:
         st.error(f"An error occurred during summarization: {str(e)}")
         return None
 
 # Page title and layout
-st.set_page_config(page_title='🦜🔗 Text Summarization App')
-st.title('🦜🔗 Text Summarization App')
+st.set_page_config(page_title='🦜🔗 RAG Text Summarization App')
+st.title('🦜🔗 RAG Text Summarization App')
 
 # Text input area for user to input text
 txt_input = st.text_area('Enter your text', '', height=200)
@@ -37,21 +36,18 @@ txt_input = st.text_area('Enter your text', '', height=200)
 # Form to accept the user's text input for summarization
 response = None
 with st.form('summarize_form', clear_on_submit=True):
-    openai_api_key = st.text_input('OpenAI API Key', type='password', disabled=not txt_input)
     submitted = st.form_submit_button('Submit')
-    if submitted and openai_api_key.startswith('sk-'):
-        with st.spinner('Summarizing...'):
-            response = generate_response(txt_input, openai_api_key)
+    if submitted and txt_input:
+        with st.spinner('Summarizing with RAG...'):
+            response = generate_response_with_rag(txt_input)
 
 # Display the response if available
 if response:
     st.info(response)
 
-# Instructions for getting an OpenAI API key
-st.subheader("Get an OpenAI API key")
-st.write("You can get your own OpenAI API key by following the instructions:")
+# Instructions for getting started with Hugging Face models
+st.subheader("Hugging Face RAG Summarization")
 st.write("""
-1. Go to [OpenAI API Keys](https://platform.openai.com/account/api-keys).
-2. Click on the `+ Create new secret key` button.
-3. Enter an identifier name (optional) and click on the `Create secret key` button.
+This app uses Hugging Face's RAG model (Retrieval-Augmented Generation) to generate summaries with relevant external context.
+RAG retrieves information from a set of documents and combines that with a generative model to produce more accurate summaries.
 """)