Update app.py

app.py

@@ -1,105 +1,143 @@
-#using codes from mistralai official cookbook
-import gradio as gr
-from llama_index.llms import MistralAI
-import numpy as np
-import PyPDF2
-import faiss
 import os
+from bs4 import BeautifulSoup
+from IPython.display import Markdown, display
+from llama_index.core import Document
+from llama_index.core import Settings
 from llama_index.core import SimpleDirectoryReader
-from llama_index.embeddings import MistralAIEmbedding
-from llama_index import ServiceContext
-from llama_index.core import VectorStoreIndex, StorageContext
-from llama_index.vector_stores.milvus import MilvusVectorStore
-import textwrap
-
-
-mistral_api_key = os.environ.get("API_KEY")
-
-cli = MistralClient(api_key = mistral_api_key)
-
-def get_text_embedding(input: str):
-    embeddings_batch_response = cli.embeddings(
-          model = "mistral-embed",
-          input = input
-      )
-    return embeddings_batch_response.data[0].embedding
-
-def rag_pdf(pdfs: list, question: str) -> str:
-    chunk_size = 4096
-    chunks = []
-    for pdf in pdfs:
-        chunks += [pdf[i:i + chunk_size] for i in range(0, len(pdf), chunk_size)]
-
-    text_embeddings = np.array([get_text_embedding(chunk) for chunk in chunks])
-    d = text_embeddings.shape[1]
-    index = faiss.IndexFlatL2(d)
-    index.add(text_embeddings)
-
-    question_embeddings = np.array([get_text_embedding(question)])
-    D, I = index.search(question_embeddings, k = 4)
-    retrieved_chunk = [chunks[i] for i in I.tolist()[0]]
-    text_retrieved = "\n\n".join(retrieved_chunk)
-    return text_retrieved
-
-def load_doc(path_list):
+from llama_index.core import StorageContext
+from llama_index.core import VectorStoreIndex
+from llama_index.readers.web import SimpleWebPageReader
+
+from llama_index.vector_stores.chroma import ChromaVectorStore
+
+import chromadb
+import re
+from llama_index.llms.gemini import Gemini
+from llama_index.embeddings.gemini import GeminiEmbedding
+
+from llama_index.core import PromptTemplate
+from llama_index.core.llms import ChatMessage
+
+import uuid
+
+api_key = os.environ.get("API_KEY")
+
+llm = Gemini(api_key=api_key, model_name="models/gemini-1.5-flash-latest")
+gemini_embedding_model = GeminiEmbedding(api_key=api_key, model_name="models/embedding-001")
+
+
+
+
+# Set Global settings
+Settings.llm = llm
+Settings.embed_model = gemini_embedding_model
+
+
+
+def extract_web(url):
+    web_documents = SimpleWebPageReader().load_data(
+        [url]
+    )
+    html_content = web_documents[0].text
+    # Parse the data.
+    soup = BeautifulSoup(html_content, 'html.parser')
+    p_tags = soup.findAll('p')
+    text_content = ""
+    for each in p_tags:
+        text_content += each.text + "\n"
+    
+    # Convert back to Document format
+    documents = [Document(text=text_content)]
+    option = "web"
+    return documents, option
+
+def extract_doc(path):
     documents = SimpleDirectoryReader(input_files=path).load_data()
-    print("Document ID:", documents[0].doc_id)
-    vector_store = MilvusVectorStore(uri="./milvus_demo.db", dim=1536, overwrite=True)
-    storage_context = StorageContext.from_defaults(vector_store=vector_store)
-    index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)
-    return index
+    option = "doc"
+    return documents, option
 
 
+def create_col(documents):
+    # Create a client and a new collection
+    db_path = f'database/{str(uuid.uuid4()[:4])}'
+    client = chromadb.PersistentClient(path=db_path)
+    chroma_collection = client.get_or_create_collection("quickstart")
+    
+    # Create a vector store
+    vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
+    
+    # Create a storage context
+    storage_context = StorageContext.from_defaults(vector_store=vector_store)
+    # Create an index from the documents and save it to the disk.
+    VectorStoreIndex.from_documents(
+        documents, storage_context=storage_context
+    )
+    return db_path
 
-def ask_mistral(message: str, history: list):
+def infer(message:str, history: list):
+    print(f'message: {message}')
+    print(f'history: {history}')
     messages = []
-    docs = message["files"]
-    for couple in history:
-        if type(couple[0]) is tuple:
-            docs += couple[0][0]
+    files_list = message["files"]
+    
+
+    for prompt,answer in history:
+        if prompt is tuple:
+            files_list += prompt[0]
         else:
-            messages.append(ChatMessage(role= "user", content = couple[0]))
-            messages.append(ChatMessage(role= "assistant", content = couple[1]))
-    if docs:
-        print(docs)
-        index = load_doc(docs)
-        query_engine = index.as_query_engine()
-        response = query_engine.query(message["text"])
-
-        full_response = ""
-        for text in response.response_gen:
-            full_response += chunk.choices[0].delta.content
-            yield full_response        
+            messages.append(ChatMessage(role= "user", content = prompt))
+            messages.append(ChatMessage(role= "assistant", content = answer))
+
             
-        
+    if files_list:
+        documents, option = extract_doc(files_list)
+    else:
+        if message["text"].startswith("http://") or message["text"].startswith("https://"):
+            documents, option = extract_doc(message["text"])
+        elif not message["text"].startswith("http://") and not message["text"].startswith("https://") and len(history) == 0:
+            gr.Error("Please input an url or upload file at first.")
+            
+
+    print(documents)
+    db_path = create_col(documents)
+    
+    # Load from disk
+    load_client = chromadb.PersistentClient(path=db_path)
+    
+    # Fetch the collection
+    chroma_collection = load_client.get_collection("quickstart")
+    
+    # Fetch the vector store
+    vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
+    
+    # Get the index from the vector store
+    index = VectorStoreIndex.from_vector_store(
+        vector_store
+    )
 
 
-        
-        pdfs_extracted = []
-        for pdf in pdfs:
-            reader = PyPDF2.PdfReader(pdf)
-            txt = ""
-            for page in reader.pages:
-                txt += page.extract_text()
-            pdfs_extracted.append(txt)
+    template = (
+        """ You are an assistant for question-answering tasks.
+    Use the following context to answer the question.
+    If you don't know the answer, just say that you don't know.
+    Use five sentences maximum and keep the answer concise.\n
+    Question: {query_str} \nContext: {context_str} \nAnswer:"""
+    )
+    llm_prompt = PromptTemplate(template)
+    print(llm_prompt)
 
-        retrieved_text = rag_pdf(pdfs_extracted, message["text"])
-        print(f'retrieved_text: {retrieved_text}')
-        messages.append(ChatMessage(role = "user", content = retrieved_text + "\n\n" + message["text"]))
-    else:
-        messages.append(ChatMessage(role = "user", content = message["text"]))
-    print(f'messages: {messages}')
+    if option == "web" and len(history) == 0:
+        response = "Get the web data! You can ask it."   
+    else: 
+        question = message['text']
+        query_engine = index.as_query_engine(text_qa_template=llm_prompt)
+        response = query_engine.query(question)
+
+    return response
     
-    full_response = ""
 
-    response = cli.chat_stream(
-        model = "open-mistral-7b",
-        messages = messages, 
-        max_tokens = 4096)
-          
-    for chunk in response:
-        full_response += chunk.choices[0].delta.content
-        yield full_response
+
+
 
 
 
@@ -108,7 +146,7 @@ chatbot = gr.Chatbot()
 with gr.Blocks(theme="soft") as demo:
     gr.ChatInterface(
         fn = ask_mistral,
-        title = "Ask Mistral and talk to your PDFs", 
+        title = "RAG demo", 
         multimodal = True,
         chatbot=chatbot,
     )