import streamlit as st
from langchain.text_splitter import CharacterTextSplitter
from langchain_community.document_loaders import TextLoader
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.chains import RetrievalQA
from langchain.embeddings.sentence_transformer import SentenceTransformerEmbeddings

def get_text():
    input_text = st.text_input("You: ", key="input")
    return input_text


user_input = get_text()
submit = st.button('Get Answer')


loader = TextLoader('India.txt')
documents =loader.load()

text_splitter = CharacterTextSplitter (chunk_size=200,
chunk_overlap=0)

texts= text_splitter.split_documents(documents)



embeddings = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
db = Chroma.from_documents(texts, embeddings)
db._collection.get(include=['embeddings'])
retriever = db.as_retriever(search_kwargs={"k": 1})

if user_input and submit:
    
    docs = retriever.get_relevant_documents(user_input)
    st.write("Answer")
    document = docs[0]
    page_content = document.page_content

    st.write(page_content)
# st.text(file_content)