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()