add jina embeddings and reranker

README.md

@@ -1,11 +1,11 @@
 ---
 title: 01aiYi NvidiaEmbed
-emoji: 💬
-colorFrom: yellow
-colorTo: purple
+emoji: ☯️🧠🛌🏻🥟🧩
+colorFrom: blue
+colorTo: red
 sdk: gradio
 sdk_version: 4.36.1
-app_file: langchainapp.py
-pinned: false
+app_file: yijinaembed.py
+pinned: true
 license: mit
 ---

app.py

@@ -1,235 +0,0 @@
-# app.py
-import spaces
-from torch.nn import DataParallel  
-from torch import Tensor  
-from transformers import AutoTokenizer, AutoModel  
-from huggingface_hub import InferenceClient  
-from openai import OpenAI  
-from langchain_community.document_loaders import UnstructuredFileLoader
-from langchain_chroma import Chroma
-from chromadb import Documents, EmbeddingFunction, Embeddings  
-from chromadb.config import Settings  
-import chromadb #import HttpClient 
-import os
-import tempfile
-import re 
-import uuid  
-import gradio as gr  
-import torch  
-import torch.nn.functional as F  
-from dotenv import load_dotenv
-from utils import load_env_variables, parse_and_route, escape_special_characters 
-from globalvars import API_BASE, intention_prompt, tasks, system_message, model_name, metadata_prompt 
-from sentence_transformers import SentenceTransformer
-
-
-load_dotenv()
-
-os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:30'  
-os.environ['CUDA_LAUNCH_BLOCKING'] = '1'  
-os.environ['CUDA_CACHE_DISABLE'] = '1'  
-  
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  
-
-# Ensure the temporary directory exists
-temp_dir = '/tmp/gradio/'
-os.makedirs(temp_dir, exist_ok=True)
-
-# Set Gradio cache directory
-gr.components.file.GRADIO_CACHE = temp_dir
-  
-### Utils  
-hf_token, yi_token = load_env_variables()  
-  
-def clear_cuda_cache():  
-    torch.cuda.empty_cache()  
-  
-client = OpenAI(api_key=yi_token, base_url=API_BASE)
-
-chroma_client = chromadb.Client(Settings())  
-  
-# Create a collection  
-chroma_collection = chroma_client.create_collection("all-my-documents")  
-
-class EmbeddingGenerator:  
-    def __init__(self, model_name: str, token: str, intention_client):  
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name, token=token, trust_remote_code=True)  
-        self.model = AutoModel.from_pretrained(model_name, token=token, trust_remote_code=True).to(self.device)  
-        self.intention_client = intention_client  
-  
-    def clear_cuda_cache(self):  
-        torch.cuda.empty_cache()  
-
-    @spaces.GPU  
-    def compute_embeddings(self, input_text: str):  
-        escaped_input_text = escape_special_characters(input_text)  
-        intention_completion = self.intention_client.chat.completions.create(  
-            model="yi-large",  
-            messages=[  
-                {"role": "system", "content": escape_special_characters(intention_prompt)},  
-                {"role": "user", "content": escaped_input_text}  
-            ]  
-        )  
-        intention_output = intention_completion.choices[0].message.content  
-        # Parse and route the intention  
-        parsed_task = parse_and_route(intention_output)  
-        selected_task = parsed_task  
-        # Construct the prompt  
-        if selected_task in tasks:
-            task_description = tasks[selected_task]  
-        else:
-            task_description = tasks["DEFAULT"]
-            print(f"Selected task not found: {selected_task}")  
-  
-        query_prefix = f"Instruct: {task_description}\nQuery: "  
-        queries = [escaped_input_text]  
-  
-        # Get the metadata  
-        metadata_completion = self.intention_client.chat.completions.create(  
-            model="yi-large",  
-            messages=[  
-                {"role": "system", "content": escape_special_characters(metadata_prompt)},  
-                {"role": "user", "content": escaped_input_text}  
-            ]  
-        )  
-        metadata_output = metadata_completion.choices[0].message.content
-        metadata = self.extract_metadata(metadata_output)  
-  
-        # Get the embeddings  
-        with torch.no_grad():  
-            inputs = self.tokenizer(queries, return_tensors='pt', padding=True, truncation=True, max_length=4096).to(self.device)  
-            outputs = self.model(**inputs)
-            query_embeddings = outputs["sentence_embeddings"].mean(dim=1)    
-            query_embeddings = outputs.last_hidden_state.mean(dim=1)  
-  
-            # Normalize embeddings  
-            query_embeddings = F.normalize(query_embeddings, p=2, dim=1)  
-            embeddings_list = query_embeddings.detach().cpu().numpy().tolist()  
-  
-            self.clear_cuda_cache()  
-            return embeddings_list, metadata   
-  
-    def extract_metadata(self, metadata_output: str):  
-        # Regex pattern to extract key-value pairs  
-        pattern = re.compile(r'\"(\w+)\": \"([^\"]+)\"')  
-        matches = pattern.findall(metadata_output)  
-        metadata = {key: value for key, value in matches}  
-        return metadata 
-
-class MyEmbeddingFunction(EmbeddingFunction):  
-    def __init__(self, model_name: str, token: str, intention_client):  
-        self.model_name = model_name  
-        self.token = token  
-        self.intention_client = intention_client  
-  
-    def create_embedding_generator(self):  
-        return EmbeddingGenerator(self.model_name, self.token, self.intention_client)  
-  
-    def __call__(self, input: Documents) -> (Embeddings, list):  
-        embedding_generator = self.create_embedding_generator()  
-        embeddings_with_metadata = [embedding_generator.compute_embeddings(doc.page_content) for doc in input]  
-        embeddings = [item[0] for item in embeddings_with_metadata]  
-        metadata = [item[1] for item in embeddings_with_metadata]  
-        embeddings_flattened = [emb for sublist in embeddings for emb in sublist]  
-        metadata_flattened = [meta for sublist in metadata for meta in sublist]  
-        return embeddings_flattened, metadata_flattened  
-  
-def load_documents(file_path: str, mode: str = "elements"):  
-    loader = UnstructuredFileLoader(file_path, mode=mode)  
-    docs = loader.load()  
-    return [doc.page_content for doc in docs]  
-
-def initialize_chroma(collection_name: str, embedding_function: MyEmbeddingFunction):  
-    db = Chroma(client=chroma_client, collection_name=collection_name, embedding_function=embedding_function)  
-    return db  
-  
-def add_documents_to_chroma(documents: list, embedding_function: MyEmbeddingFunction):  
-    for doc in documents:  
-        embeddings, metadata = embedding_function.create_embedding_generator().compute_embeddings(doc)  
-        for embedding, meta in zip(embeddings, metadata):  
-            chroma_collection.add(  
-                ids=[str(uuid.uuid1())],  
-                documents=[doc],  
-                embeddings=[embedding],  
-                metadatas=[meta]  
-            )  
-  
-def query_chroma(query_text: str, embedding_function: MyEmbeddingFunction):  
-    query_embeddings, query_metadata = embedding_function.create_embedding_generator().compute_embeddings(query_text)  
-    result_docs = chroma_collection.query(  
-        query_texts=[query_text],  
-        n_results=2  
-    )  
-    return result_docs  
-
-# Initialize clients  
-intention_client = OpenAI(api_key=yi_token, base_url=API_BASE)  
-embedding_generator = EmbeddingGenerator(model_name=model_name, token=hf_token, intention_client=intention_client)  
-embedding_function = MyEmbeddingFunction(model_name=model_name, token=hf_token, intention_client=intention_client)  
-chroma_db = initialize_chroma(collection_name="Tonic-instruct", embedding_function=embedding_function)
-
-def respond(  
-    message,  
-    history: list[tuple[str, str]],  
-    system_message,  
-    max_tokens,  
-    temperature,  
-    top_p,  
-):  
-    retrieved_text = query_documents(message)  
-    messages = [{"role": "system", "content": escape_special_characters(system_message)}]  
-    for val in history:  
-        if val[0]:  
-            messages.append({"role": "user", "content": val[0]})  
-        if val[1]:  
-            messages.append({"role": "assistant", "content": val[1]})  
-    messages.append({"role": "user", "content": f"{retrieved_text}\n\n{escape_special_characters(message)}"})  
-    response = ""  
-    for message in intention_client.chat_completion(  
-        messages,  
-        max_tokens=max_tokens,  
-        stream=True,  
-        temperature=temperature,  
-        top_p=top_p,  
-    ):  
-        token = message.choices[0].delta.content  
-        response += token  
-        yield response  
-  
-def upload_documents(files):  
-    for file in files:  
-        loader = UnstructuredFileLoader(file.name)  
-        documents = loader.load()  
-        add_documents_to_chroma(documents, embedding_function)  
-    return "Documents uploaded and processed successfully!"  
-  
-def query_documents(query):  
-    results = query_chroma(query, embedding_function)  
-    return "\n\n".join([result.content for result in results])  
-  
-with gr.Blocks() as demo:  
-    with gr.Tab("Upload Documents"):  
-        document_upload = gr.File(file_count="multiple", file_types=["document"])  
-        upload_button = gr.Button("Upload and Process")  
-        upload_button.click(upload_documents, inputs=document_upload, outputs=gr.Text())  
-  
-    with gr.Tab("Ask Questions"):  
-        with gr.Row():  
-            chat_interface = gr.ChatInterface(  
-                respond,  
-                additional_inputs=[  
-                    gr.Textbox(value="You are a friendly Chatbot.", label="System message"),  
-                    gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),  
-                    gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),  
-                    gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),  
-                ],  
-            )  
-            query_input = gr.Textbox(label="Query")  
-            query_button = gr.Button("Query")  
-            query_output = gr.Textbox()  
-            query_button.click(query_documents, inputs=query_input, outputs=query_output)  
-  
-if __name__ == "__main__":  
-    # os.system("chroma run --host localhost --port 8000 &")
-    demo.launch()

globalvars.py

@@ -3,7 +3,7 @@
 API_BASE = "https://api.01.ai/v1"
 API_KEY = "your key"
 
-model_name = 'nvidia/NV-Embed-v1'
+model_name = "jinaai/jina-embeddings-v3"
 
 title = """
 # 👋🏻Welcome to 🙋🏻‍♂️Tonic's 📽️Nvidia 🛌🏻Embed V-1 !"""
@@ -15,76 +15,47 @@ Join us : 🌟TeamTonic🌟 is always making cool demos! Join our active builder
 """
 
 tasks = {
-        'ClimateFEVER': 'Given a claim about climate change, retrieve documents that support or refute the claim',
-        'DBPedia': 'Given a query, retrieve relevant entity descriptions from DBPedia',
-        'FEVER': 'Given a claim, retrieve documents that support or refute the claim',
-        'FiQA2018': 'Given a financial question, retrieve user replies that best answer the question',
-        'HotpotQA': 'Given a multi-hop question, retrieve documents that can help answer the question',
-        'MSMARCO': 'Given a web search query, retrieve relevant passages that answer the query',
-        'NFCorpus': 'Given a question, retrieve relevant documents that best answer the question',
-        'NQ': 'Given a question, retrieve Wikipedia passages that answer the question',
-        'QuoraRetrieval': 'Given a question, retrieve questions that are semantically equivalent to the given question',
-        'SCIDOCS': 'Given a scientific paper title, retrieve paper abstracts that are cited by the given paper',
-        'DEFAULT': 'Given a query, retrieve relevant entity descriptions from DBPedia',
+    'retrieval.query': 'Used for query embeddings in asymmetric retrieval tasks',
+    'retrieval.passage': 'Used for passage embeddings in asymmetric retrieval tasks',
+    'separation': 'Used for embeddings in clustering and re-ranking applications',
+    'classification': 'Used for embeddings in classification tasks',
+    'text-matching': 'Used for embeddings in tasks that quantify similarity between two texts, such as STS or symmetric retrieval tasks',
+    'DEFAULT': 'Used for general-purpose embeddings when no specific task is specified'
 }
 
-intention_prompt= """
+intention_prompt = """
+{
   "type": "object",
   "properties": {
-    "ClimateFEVER": {
+    "retrieval.query": {
       "type": "boolean",
-      "description" : "select this for climate science related text"
+      "description": "Select this for query embeddings in asymmetric retrieval tasks"
     },
-    "DBPedia": {
+    "retrieval.passage": {
       "type": "boolean",
-      "description" : "select this for encyclopedic related knowledge"
+      "description": "Select this for passage embeddings in asymmetric retrieval tasks"
     },
-    "FEVER": {
+    "separation": {
       "type": "boolean",
-      "description": "select this to verify a claim or embed a claim"
+      "description": "Select this for embeddings in clustering and re-ranking applications"
     },
-    "FiQA2018": {
+    "classification": {
       "type": "boolean",
-      "description" : "select this for financial questions or topics"
+      "description": "Select this for embeddings in classification tasks"
     },
-    "HotpotQA": {
+    "text-matching": {
       "type": "boolean",
-      "description" : "select this for a multi-hop question or for texts that provide multihop claims"
-    },
-    "MSMARCO": {
-      "type": "boolean",
-      "description": "Given a web search query, retrieve relevant passages that answer the query"
-    },
-    "NFCorpus": {
-      "type": "boolean",
-      "description" : "Given a question, retrieve relevant documents that best answer the question"
-    },
-    "NQ": {
-      "type": "boolean",
-      "description" : "Given a question, retrieve Wikipedia passages that answer the question"
-    },
-    "QuoraRetrieval": {
-      "type": "boolean",
-      "description": "Given a question, retrieve questions that are semantically equivalent to the given question"
-    },
-    "SCIDOCS": {
-      "type": "boolean",
-      "description": "Given a scientific paper title, retrieve paper abstracts that are cited by the given paper"
+      "description": "Select this for embeddings in tasks that quantify similarity between two texts, such as STS or symmetric retrieval tasks"
     }
   },
   "required": [
-    "ClimateFEVER",
-    "DBPedia",
-    "FEVER",
-    "FiQA2018",
-    "HotpotQA",
-    "MSMARCO",
-    "NFCorpus",
-    "NQ",
-    "QuoraRetrieval",
-    "SCIDOCS",
+    "retrieval.query",
+    "retrieval.passage",
+    "separation",
+    "classification",
+    "text-matching"
   ]
-produce a complete json schema."
+}
 
 you will recieve a text , classify the text according to the schema above. ONLY PROVIDE THE FINAL JSON , DO NOT PRODUCE ANY ADDITION INSTRUCTION :"""
 

langchainapp.py

@@ -1,243 +0,0 @@
-# app.py  
-import spaces
-from torch.nn import DataParallel  
-from torch import Tensor  
-from transformers import AutoTokenizer, AutoModel  
-from huggingface_hub import InferenceClient  
-from openai import OpenAI  
-from langchain_community.embeddings import HuggingFaceInstructEmbeddings  
-from langchain_community.document_loaders import UnstructuredFileLoader
-from langchain_chroma import Chroma
-from chromadb import Documents, EmbeddingFunction, Embeddings  
-from chromadb.config import Settings  
-import chromadb #import HttpClient 
-from typing import List, Tuple, Dict, Any  
-import os  
-import re 
-import uuid  
-import gradio as gr  
-import torch  
-import torch.nn.functional as F  
-from dotenv import load_dotenv
-from utils import load_env_variables, parse_and_route , escape_special_characters 
-from globalvars import API_BASE, intention_prompt, tasks, system_message, model_name , metadata_prompt 
-# import time  
-# import httpx  
-
-from langchain_community.chat_models import ChatOpenAI
-from langchain.retrievers.document_compressors import LLMChainExtractor
-from langchain.retrievers.multi_query import MultiQueryRetriever
-from langchain.retrievers import ContextualCompressionRetriever
-from langchain.prompts.chat import ChatPromptTemplate, HumanMessagePromptTemplate
-# from langchain.vectorstores import Chroma
-  
-
-
-load_dotenv()
-
-os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:50'  
-os.environ['CUDA_LAUNCH_BLOCKING'] = '1'  
-os.environ['CUDA_CACHE_DISABLE'] = '1'  
-  
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  
-### Utils  
-hf_token, yi_token = load_env_variables()  
-
-# tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token, trust_remote_code=True)  
-# Lazy load model  
-model = None  
-  
-@spaces.GPU  
-def load_model():  
-    global model  
-    if model is None:  
-        from transformers import AutoModel  
-        model = AutoModel.from_pretrained(model_name, token=hf_token, trust_remote_code=True).to(device)  
-    return model  
-  
-# Load model  
-nvidiamodel = load_model()   
-# nvidiamodel.set_pooling_include_prompt(include_prompt=False) 
-
-def clear_cuda_cache():  
-    torch.cuda.empty_cache()  
-  
-client = OpenAI(api_key=yi_token, base_url=API_BASE)
-
-chroma_client = chromadb.Client(Settings())  
-  
-# Create a collection  
-chroma_collection = chroma_client.create_collection("all-my-documents")  
-
-@spaces.GPU  
-class MyEmbeddingFunction(EmbeddingFunction):  
-    def __init__(self, model_name: str, token: str, intention_client):  
-        self.model_name = model_name  
-        self.token = token  
-        self.intention_client = intention_client  
-        self.hf_embeddings = HuggingFaceInstructEmbeddings(  
-            model_name=model_name,  
-            model_kwargs={'device': 'cuda' if torch.cuda.is_available() else 'cpu'},  
-            encode_kwargs={'normalize_embeddings': True}  
-        )  
-  
-    def create_embedding_generator(self):  
-        return self.hf_embeddings  
-  
-    def __call__(self, input: Documents) -> (List[List[float]], List[Dict[str, Any]]):  
-        embeddings_with_metadata = [self.compute_embeddings(doc.page_content) for doc in input]  
-        embeddings = [item[0] for item in embeddings_with_metadata]  
-        metadata = [item[1] for item in embeddings_with_metadata]  
-        embeddings_flattened = [emb for sublist in embeddings for emb in sublist]  
-        metadata_flattened = [meta for sublist in metadata for meta in sublist]  
-        return embeddings_flattened, metadata_flattened  
-    
-    @spaces.GPU  
-    def compute_embeddings(self, input_text: str):  
-        escaped_input_text = escape_special_characters(input_text)  
-          
-        # Get the intention  
-        intention_completion = self.intention_client.chat.completions.create(  
-            model="yi-large",  
-            messages=[  
-                {"role": "system", "content": escape_special_characters(intention_prompt)},  
-                {"role": "user", "content": escaped_input_text}  
-            ]  
-        )  
-        intention_output = intention_completion.choices[0].message.content  
-        parsed_task = parse_and_route(intention_output)  
-        selected_task = parsed_task if parsed_task in tasks else "DEFAULT"  
-        task_description = tasks[selected_task]  
-        # query_prefix = "Instruct: " +tasks[selected_task] +"\nQuery: "  
-        # Construct the embed_instruction and query_instruction dynamically  
-        embed_instruction = f"Instruct: {task_description}" +"\nQuery:"  
-        # query_instruction = f""  
-          
-        # Update the hf_embeddings object with the new instructions  
-        self.hf_embeddings.embed_instruction = embed_instruction  
-        # self.hf_embeddings.query_instruction = query_instruction  
-          
-        # Get the metadata  
-        metadata_completion = self.intention_client.chat.completions.create(  
-            model="yi-large",  
-            messages=[  
-                {"role": "system", "content": escape_special_characters(metadata_prompt)},  
-                {"role": "user", "content": escaped_input_text}  
-            ]  
-        )  
-        metadata_output = metadata_completion.choices[0].message.content  
-        metadata = self.extract_metadata(metadata_output)  
-          
-        # Get the embeddings  
-        embeddings = self.hf_embeddings.embed_documents([escaped_input_text])  
-        return embeddings[0], metadata  
-  
-    def extract_metadata(self, metadata_output: str) -> Dict[str, str]:  
-        pattern = re.compile(r'\"(\w+)\": \"([^\"]+)\"')  
-        matches = pattern.findall(metadata_output)  
-        metadata = {key: value for key, value in matches}  
-        return metadata  
-  
-def load_documents(file_path: str, mode: str = "elements"):  
-    loader = UnstructuredFileLoader(file_path, mode=mode)  
-    docs = loader.load()  
-    return [doc.page_content for doc in docs]  
-
-def initialize_chroma(collection_name: str, embedding_function: MyEmbeddingFunction):  
-    db = Chroma(client=chroma_client, collection_name=collection_name, embedding_function=embedding_function)  
-    return db  
-  
-def add_documents_to_chroma(documents: list, embedding_function: MyEmbeddingFunction):  
-    for doc in documents:  
-        embeddings, metadata = embedding_function.compute_embeddings(doc)  
-        for embedding, meta in zip(embeddings, metadata):  
-            chroma_collection.add(  
-                ids=[str(uuid.uuid1())],  
-                documents=[doc],  
-                embeddings=[embedding],  
-                metadatas=[meta]  
-            )  
-  
-def query_chroma(query_text: str, embedding_function: MyEmbeddingFunction):  
-    model = load_model()  
-    query_embeddings, query_metadata = embedding_function.compute_embeddings(query_text)  
-    result_docs = chroma_collection.query(  
-        query_texts=[query_text],  
-        n_results=3  
-    )  
-    return result_docs 
-
-
-def answer_query(message: str, chat_history: List[Tuple[str, str]]):  
-    base_compressor = LLMChainExtractor.from_llm(intention_client)  
-    db = Chroma(persist_directory="output/general_knowledge", embedding_function=embedding_function)  
-    base_retriever = db.as_retriever()  
-    mq_retriever = MultiQueryRetriever.from_llm(retriever=base_retriever, llm=intention_client)  
-    compression_retriever = ContextualCompressionRetriever(base_compressor=base_compressor, base_retriever=mq_retriever)  
-  
-    matched_docs = compression_retriever.get_relevant_documents(query=message)  
-    context = ""  
-    for doc in matched_docs:  
-        page_content = doc.page_content  
-        context += page_content  
-        context += "\n\n"  
-  
-    template = """  
-    Answer the following question only by using the context given below in the triple backticks, do not use any other information to answer the question.  
-    If you can't answer the given question with the given context, you can return an empty string ('')  
-    Context: ```{context}```  
-    ----------------------------  
-    Question: {query}  
-    ----------------------------  
-    Answer: """  
-  
-    human_message_prompt = HumanMessagePromptTemplate.from_template(template=template)  
-    chat_prompt = ChatPromptTemplate.from_messages([human_message_prompt])  
-    prompt = chat_prompt.format_prompt(query=message, context=context)  
-    response = intention_client.chat(messages=prompt.to_messages()).content  
-    chat_history.append((message, response))  
-    return "", chat_history  
-  
-
-# Initialize clients  
-intention_client = OpenAI(api_key=yi_token, base_url=API_BASE)  
-embedding_function = MyEmbeddingFunction(model_name=model_name, token=hf_token, intention_client=intention_client)  
-chroma_db = initialize_chroma(collection_name="Tonic-instruct", embedding_function=embedding_function) 
-  
-def upload_documents(files):  
-    for file in files:  
-        loader = UnstructuredFileLoader(file.name)  
-        documents = loader.load()  
-        add_documents_to_chroma(documents, embedding_function)  
-    return "Documents uploaded and processed successfully!"  
-  
-def query_documents(query):  
-    model = load_model()  
-    results = query_chroma(query)  
-    return "\n\n".join([result.content for result in results])  
-  
-with gr.Blocks() as demo:  
-    with gr.Tab("Upload Documents"):  
-        document_upload = gr.File(file_count="multiple", file_types=["document"])  
-        upload_button = gr.Button("Upload and Process")  
-        upload_button.click(upload_documents, inputs=document_upload, outputs=gr.Text())  
-  
-    with gr.Tab("Ask Questions"):  
-        with gr.Row():  
-            chat_interface = gr.ChatInterface(  
-                answer_query,  
-                additional_inputs=[  
-                    gr.Textbox(value="You are a friendly Chatbot.", label="System message"),  
-                    gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),  
-                    gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),  
-                    gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),  
-                ],  
-            )  
-            query_input = gr.Textbox(label="Query")  
-            query_button = gr.Button("Query")  
-            query_output = gr.Textbox()  
-            query_button.click(query_documents, inputs=query_input, outputs=query_output)  
-  
-if __name__ == "__main__":  
-    # os.system("chroma run --host localhost --port 8000 &")
-    demo.launch()  

yijinaembed.py

@@ -0,0 +1,231 @@
+# app.py
+import os
+import re
+import uuid
+import gradio as gr
+import torch
+import torch.nn.functional as F
+from dotenv import load_dotenv
+from typing import List, Tuple, Dict, Any
+from transformers import AutoTokenizer, AutoModel
+from openai import OpenAI
+from langchain_community.document_loaders import UnstructuredFileLoader
+from langchain_chroma import Chroma
+from chromadb import Documents, EmbeddingFunction, Embeddings
+from chromadb.config import Settings
+import chromadb
+from utils import load_env_variables, parse_and_route, escape_special_characters
+from globalvars import API_BASE, intention_prompt, tasks, system_message, metadata_prompt, model_name
+import spaces
+from langchain.retrievers.contextual_compression import ContextualCompressionRetriever
+from langchain_community.document_compressors.jina_rerank import JinaRerank
+from langchain import hub
+from langchain.chains import create_retrieval_chain
+from langchain.chains.retrieval import create_stuff_documents_chain
+
+load_dotenv()
+
+# os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:180'
+# os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
+# os.environ['CUDA_CACHE_DISABLE'] = '1'
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+hf_token, yi_token = load_env_variables()
+
+tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token, trust_remote_code=True)
+model = None
+
+@spaces.GPU
+def load_model():
+    global model
+    if model is None:
+        model = AutoModel.from_pretrained(model_name, token=hf_token, trust_remote_code=True).to(device)
+    return model
+
+# Load model
+jina_model = load_model()
+
+def clear_cuda_cache():
+    torch.cuda.empty_cache()
+
+client = OpenAI(api_key=yi_token, base_url=API_BASE)
+
+chroma_client = chromadb.Client(Settings())
+
+chroma_collection = chroma_client.create_collection("all-my-documents")
+
+class JinaEmbeddingFunction(EmbeddingFunction):
+    def __init__(self, model, tokenizer, intention_client):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.intention_client = intention_client
+
+    def __call__(self, input: Documents) -> Tuple[List[List[float]], List[Dict[str, Any]]]:
+        embeddings_with_metadata = [self.compute_embeddings(doc) for doc in input]
+        embeddings = [item[0] for item in embeddings_with_metadata]
+        metadata = [item[1] for item in embeddings_with_metadata]
+        return embeddings, metadata
+
+    @spaces.GPU
+    def compute_embeddings(self, input_text: str):
+        escaped_input_text = escape_special_characters(input_text)
+
+        # Get the intention
+        intention_completion = self.intention_client.chat.completions.create(
+            model="yi-large",
+            messages=[
+                {"role": "system", "content": escape_special_characters(intention_prompt)},
+                {"role": "user", "content": escaped_input_text}
+            ]
+        )
+        intention_output = intention_completion.choices[0].message.content
+        parsed_task = parse_and_route(intention_output)
+        selected_task = parsed_task if parsed_task in tasks else "DEFAULT"
+        task = tasks[selected_task]
+
+        # Get the metadata
+        metadata_completion = self.intention_client.chat.completions.create(
+            model="yi-large",
+            messages=[
+                {"role": "system", "content": escape_special_characters(metadata_prompt)},
+                {"role": "user", "content": escaped_input_text}
+            ]
+        )
+        metadata_output = metadata_completion.choices[0].message.content
+        metadata = self.extract_metadata(metadata_output)
+
+        # Compute embeddings using Jina model
+        encoded_input = self.tokenizer(escaped_input_text, padding=True, truncation=True, return_tensors="pt").to(device)
+        with torch.no_grad():
+            model_output = self.model(**encoded_input, task=task)
+
+        embeddings = self.mean_pooling(model_output, encoded_input["attention_mask"])
+        embeddings = F.normalize(embeddings, p=2, dim=1)
+
+        return embeddings.cpu().numpy().tolist()[0], metadata
+
+    def extract_metadata(self, metadata_output: str) -> Dict[str, str]:
+        pattern = re.compile(r'\"(\w+)\": \"([^\"]+)\"')
+        matches = pattern.findall(metadata_output)
+        metadata = {key: value for key, value in matches}
+        return metadata
+
+    @staticmethod
+    def mean_pooling(model_output, attention_mask):
+        token_embeddings = model_output[0]
+        input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
+        return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
+
+def load_documents(file_path: str, mode: str = "elements"):
+    loader = UnstructuredFileLoader(file_path, mode=mode)
+    docs = loader.load()
+    return [doc.page_content for doc in docs]
+
+def initialize_chroma(collection_name: str, embedding_function: JinaEmbeddingFunction):
+    db = Chroma(client=chroma_client, collection_name=collection_name, embedding_function=embedding_function)
+    return db
+
+@spaces.GPU
+def add_documents_to_chroma(documents: list, embedding_function: JinaEmbeddingFunction):
+    for doc in documents:
+        embeddings, metadata = embedding_function.compute_embeddings(doc)
+        chroma_collection.add(
+            ids=[str(uuid.uuid1())],
+            documents=[doc],
+            embeddings=[embeddings],
+            metadatas=[metadata]
+        )
+
+@spaces.GPU
+def rerank_documents(query: str, documents: List[str]) -> List[str]:
+    compressor = JinaRerank()
+    retriever = chroma_db.as_retriever(search_kwargs={"k": 20})
+    compression_retriever = ContextualCompressionRetriever(
+        base_compressor=compressor, base_retriever=retriever
+    )
+    
+    compressed_docs = compression_retriever.get_relevant_documents(query)
+    
+    return [doc.page_content for doc in compressed_docs]
+
+def query_chroma(query_text: str, embedding_function: JinaEmbeddingFunction):
+    query_embeddings, query_metadata = embedding_function.compute_embeddings(query_text)
+    result_docs = chroma_collection.query(
+        query_embeddings=[query_embeddings],
+        n_results=3
+    )
+    return result_docs
+
+@spaces.GPU
+def answer_query(message: str, chat_history: List[Tuple[str, str]], system_message: str, max_new_tokens: int, temperature: float, top_p: float):
+    # Query Chroma for relevant documents
+    results = query_chroma(message, embedding_function)
+    context = "\n\n".join([result['document'] for result in results['documents'][0]])
+
+    # Rerank the documents
+    reranked_docs = rerank_documents(message, context.split("\n\n"))
+    reranked_context = "\n\n".join(reranked_docs)
+
+    # Prepare the prompt for YI model
+    prompt = f"{system_message}\n\nContext: {reranked_context}\n\nHuman: {message}\n\nAssistant:"
+
+    # Generate response using YI model
+    response = client.chat.completions.create(
+        model="yi-large",
+        messages=[
+            {"role": "system", "content": system_message},
+            {"role": "user", "content": f"Context: {reranked_context}\n\nHuman: {message}"}
+        ],
+        max_tokens=max_new_tokens,
+        temperature=temperature,
+        top_p=top_p
+    )
+
+    assistant_response = response.choices[0].message.content
+    chat_history.append((message, assistant_response))
+    return "", chat_history
+
+# Initialize clients
+intention_client = OpenAI(api_key=yi_token, base_url=API_BASE)
+embedding_function = JinaEmbeddingFunction(jina_model, tokenizer, intention_client)
+chroma_db = initialize_chroma(collection_name="Jina-embeddings", embedding_function=embedding_function)
+
+@spaces.GPU
+def upload_documents(files):
+    for file in files:
+        loader = UnstructuredFileLoader(file.name)
+        documents = loader.load()
+        add_documents_to_chroma([doc.page_content for doc in documents], embedding_function)
+    return "Documents uploaded and processed successfully!"
+
+@spaces.GPU
+def query_documents(query):
+    results = query_chroma(query, embedding_function)
+    reranked_docs = rerank_documents(query, [result for result in results['documents'][0]])
+    return "\n\n".join(reranked_docs)
+
+with gr.Blocks() as demo:
+    with gr.Tab("Upload Documents"):
+        document_upload = gr.File(file_count="multiple", file_types=["document"])
+        upload_button = gr.Button("Upload and Process")
+        upload_button.click(upload_documents, inputs=document_upload, outputs=gr.Text())
+
+    with gr.Tab("Ask Questions"):
+        with gr.Row():
+            chat_interface = gr.ChatInterface(
+                answer_query,
+                additional_inputs=[
+                    gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
+                    gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
+                    gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
+                    gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
+                ],
+            )
+            query_input = gr.Textbox(label="Query")
+            query_button = gr.Button("Query")
+            query_output = gr.Textbox()
+            query_button.click(query_documents, inputs=query_input, outputs=query_output)
+
+if __name__ == "__main__":
+    demo.launch()