Create app.py

app.py

@@ -0,0 +1,178 @@
+model_name = "qwen:0.5b-chat"
+
+import os 
+
+os.system("sudo apt install lshw")
+os.system("curl https://ollama.ai/install.sh | sh")
+
+import nest_asyncio
+nest_asyncio.apply()
+
+import os
+import asyncio
+
+# Run Async Ollama
+# Taken from: https://stackoverflow.com/questions/77697302/how-to-run-ollama-in-google-colab
+# NB: You may need to set these depending and get cuda working depending which backend you are running.
+# Set environment variable for NVIDIA library
+# Set environment variables for CUDA
+os.environ['PATH'] += ':/usr/local/cuda/bin'
+# Set LD_LIBRARY_PATH to include both /usr/lib64-nvidia and CUDA lib directories
+os.environ['LD_LIBRARY_PATH'] = '/usr/lib64-nvidia:/usr/local/cuda/lib64'
+
+async def run_process(cmd):
+    print('>>> starting', *cmd)
+    process = await asyncio.create_subprocess_exec(
+        *cmd,
+        stdout=asyncio.subprocess.PIPE,
+        stderr=asyncio.subprocess.PIPE
+    )
+
+    # define an async pipe function
+    async def pipe(lines):
+        async for line in lines:
+            print(line.decode().strip())
+
+        await asyncio.gather(
+            pipe(process.stdout),
+            pipe(process.stderr),
+        )
+
+    # call it
+    await asyncio.gather(pipe(process.stdout), pipe(process.stderr))
+
+import asyncio
+import threading
+
+async def start_ollama_serve():
+    await run_process(['ollama', 'serve'])
+
+def run_async_in_thread(loop, coro):
+    asyncio.set_event_loop(loop)
+    loop.run_until_complete(coro)
+    loop.close()
+
+# Create a new event loop that will run in a new thread
+new_loop = asyncio.new_event_loop()
+
+# Start ollama serve in a separate thread so the cell won't block execution
+thread = threading.Thread(target=run_async_in_thread, args=(new_loop, start_ollama_serve()))
+thread.start()
+
+# Load up model
+
+os.system(f"ollama pull {model_name}")
+
+# Download Data
+
+os.system("wget -O data.txt https://drive.google.com/uc?id=1uMvEYq17LsvTkX8bU5Fq-2FcG16XbrAW")
+
+from llama_index import SimpleDirectoryReader
+from llama_index import Document
+from llama_index.embeddings import HuggingFaceEmbedding
+from llama_index import (
+    SimpleDirectoryReader,
+    VectorStoreIndex,
+    ServiceContext,
+)
+from llama_index.llms import Ollama
+from llama_index import ServiceContext, VectorStoreIndex, StorageContext
+from llama_index.indices.postprocessor import SentenceTransformerRerank
+from llama_index import load_index_from_storage
+from llama_index.node_parser import HierarchicalNodeParser
+
+from llama_index.node_parser import get_leaf_nodes
+from llama_index import StorageContext
+from llama_index.retrievers import AutoMergingRetriever
+from llama_index.indices.postprocessor import SentenceTransformerRerank
+from llama_index.query_engine import RetrieverQueryEngine
+import gradio as gr
+import os
+from llama_index import get_response_synthesizer
+from llama_index.chat_engine.condense_question import (
+    CondenseQuestionChatEngine,
+)
+
+from llama_index import set_global_service_context
+
+def build_automerging_index(
+    documents,
+    llm,
+    embed_model,
+    save_dir="merging_index",
+    chunk_sizes=None,
+):
+    chunk_sizes = chunk_sizes or [2048, 512, 128]
+    node_parser = HierarchicalNodeParser.from_defaults(chunk_sizes=chunk_sizes)
+    nodes = node_parser.get_nodes_from_documents(documents)
+    leaf_nodes = get_leaf_nodes(nodes)
+    merging_context = ServiceContext.from_defaults(
+        llm=llm,
+        embed_model=embed_model,
+    )
+    set_global_service_context(merging_context)
+    storage_context = StorageContext.from_defaults()
+    storage_context.docstore.add_documents(nodes)
+
+    if not os.path.exists(save_dir):
+        automerging_index = VectorStoreIndex(
+            leaf_nodes, storage_context=storage_context, service_context=merging_context
+        )
+        automerging_index.storage_context.persist(persist_dir=save_dir)
+    else:
+        automerging_index = load_index_from_storage(
+            StorageContext.from_defaults(persist_dir=save_dir),
+            service_context=merging_context,
+        )
+    return automerging_index
+
+
+def get_automerging_query_engine(
+    automerging_index,
+    similarity_top_k=5,
+    rerank_top_n=2,
+):
+    base_retriever = automerging_index.as_retriever(similarity_top_k=similarity_top_k)
+    retriever = AutoMergingRetriever(
+        base_retriever, automerging_index.storage_context, verbose=True
+    )
+    rerank = SentenceTransformerRerank(
+        top_n=rerank_top_n, model="BAAI/bge-reranker-base"
+    )
+    synth = get_response_synthesizer(streaming=True)
+    auto_merging_engine = RetrieverQueryEngine.from_args(
+        retriever, node_postprocessors=[rerank],response_synthesizer=synth
+    )
+    return auto_merging_engine
+
+
+llm = Ollama(model=model_name, request_timeout=300.0)
+embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-base-en-v1.5")
+
+documents = SimpleDirectoryReader(
+    input_files=["data.txt"]
+).load_data()
+
+automerging_index = build_automerging_index(
+    documents,
+    llm,
+    embed_model=embed_model,
+    save_dir="merging_index"
+)
+
+automerging_query_engine = get_automerging_query_engine(
+    automerging_index,
+)
+automerging_chat_engine = CondenseQuestionChatEngine.from_defaults(
+    query_engine=automerging_query_engine,
+)
+
+def chat(message, history):
+    res = automerging_chat_engine.stream_chat(message)
+    response = ""
+    for text in res.response_gen:
+        response+=text
+        yield response
+
+demo = gr.ChatInterface(chat)
+demo.launch()