Upload 6 files

Uploading local files in order to get this space started

generator.py

@@ -0,0 +1,62 @@
+from retriever import Retriever
+from reranker import Reranker
+from anthropic import Anthropic
+from typing import List
+
+
+retriever = Retriever()
+reranker = Reranker()
+
+
+class RAGAgent:
+    def __init__(
+            self,
+            retriever=retriever,
+            reranker=reranker,
+            anthropic_api_key: str = "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",
+            model: str = "claude-3-5-sonnet-20241022",
+            max_tokens: int = 1024,
+            temperature: float = 0.0,
+    ):
+        self.retriever = retriever
+        self.reranker = reranker
+        self.client = Anthropic(api_key=anthropic_api_key)
+        self.model = model
+        self.max_tokens = max_tokens
+        self.temperature = temperature
+
+    def get_context(self, query: str) -> List[str]:
+        # Get initial candidates from retriever
+        retrieved_docs = self.retriever.search_similar(query)
+
+        # Rerank the candidates
+        context = self.reranker.rerank(query, retrieved_docs)
+
+        return context
+
+    def generate_prompt(self, context: List[str]) -> str:
+        context = "\n".join(context)
+        prompt = f"""
+                    "אתה רופא שיניים, דובר עברית בלבד. קוראים לך 'רופא השיניים העברי האלקטרוני הראשון'. ענה למטופל על השאלה שלו על סמך הקונטקס הבא: {context}. הוסף כמה שיותר פרטים, ודאג שהתחביר יהיה תקין ויפה. תעצור כשאתה מרגיש שמיצית את עצמך. אל תמציא דברים. ואל תענה בשפות שהן לא עברית.
+                 """
+        return prompt
+
+    def get_response(self, question: str) -> str:
+        # Get relevant context
+        context = self.get_context(question)
+
+        # Generate prompt with context
+        prompt = self.generate_prompt(context)
+
+        # Get response from Claude
+        response = self.client.messages.create(
+            model=self.model,
+            max_tokens=self.max_tokens,
+            temperature=self.temperature,
+            messages=[
+                        {"role": "assistant", "content": prompt},
+                        {"role": "user", "content": f"{question}"}
+                    ]
+        )
+
+        return response.content[0].text

main.py

@@ -0,0 +1,117 @@
+import gradio as gr
+import time
+
+
+from rag_agent import RAGAgent
+
+
+rag_agent = RAGAgent()
+
+
+class ChatBot:
+    def __init__(self, rag_agent):
+        self.message_history = []
+        self.rag_agent = rag_agent
+
+    def get_response(self, message):
+        return self.rag_agent.get_response(message)
+
+    def chat(self, message, history):
+        time.sleep(1)
+        bot_response = self.get_response(message)
+        self.message_history.append((message, bot_response))
+        return bot_response
+
+
+def create_chat_interface(rag_agent=rag_agent):
+    chatbot = ChatBot(rag_agent=rag_agent)
+
+    custom_css = """
+    #chatbot {
+        direction: rtl;
+        height: 400px;
+    }
+    .message {
+        font-size: 16px;
+        text-align: right;
+    }
+    .message-wrap {
+        direction: rtl !important;
+    }
+    .message-wrap > div {
+        direction: rtl !important;
+        text-align: right !important;
+    }
+    .input-box {
+        direction: rtl !important;
+        text-align: right !important;
+    }
+    .container {
+        direction: rtl;
+    }
+    .contain {
+        direction: rtl !important;
+    }
+    .bubble {
+        direction: rtl !important;
+        text-align: right !important;
+    }
+    textarea, input {
+        direction: rtl !important;
+        text-align: right !important;
+    }
+    .user-message, .bot-message {
+        direction: rtl !important;
+        text-align: right !important;
+    }
+    """
+
+    with gr.Blocks(css=custom_css) as interface:
+        with gr.Column(elem_classes="container"):
+            gr.Markdown("רופא שיניים אלקטרוני", rtl=True)
+
+            chatbot_component = gr.Chatbot(
+                [],
+                elem_id="chatbot",
+                height=400,
+                rtl=True,
+                elem_classes="message-wrap"
+            )
+
+            with gr.Row():
+                submit_btn = gr.Button("שלח", variant="primary")
+                txt = gr.Textbox(
+                    show_label=False,
+                    placeholder="הקלד את ההודעה שלך כאן...",
+                    container=False,
+                    elem_classes="input-box",
+                    rtl=True
+                )
+
+            clear_btn = gr.Button("נקה צ'אט")
+
+        def user_message(user_message, history):
+            return "", history + [[user_message, None]]
+
+        def bot_message(history):
+            user_message = history[-1][0]
+            bot_response = chatbot.chat(user_message, history)
+            history[-1][1] = bot_response
+            return history
+
+        txt_msg = txt.submit(user_message, [txt, chatbot_component], [txt, chatbot_component], queue=False).then(
+            bot_message, chatbot_component, chatbot_component
+        )
+
+        submit_btn.click(user_message, [txt, chatbot_component], [txt, chatbot_component], queue=False).then(
+            bot_message, chatbot_component, chatbot_component
+        )
+
+        clear_btn.click(lambda: None, None, chatbot_component, queue=False)
+
+    return interface
+
+
+# Launch the interface
+chat_interface = create_chat_interface(rag_agent=rag_agent)
+chat_interface.launch(share=True)

rag_agent.py

@@ -0,0 +1,64 @@
+from anthropic import Anthropic
+from typing import List
+
+
+from retriever import Retriever
+from reranker import Reranker
+
+
+retriever = Retriever()
+reranker = Reranker()
+
+
+class RAGAgent:
+    def __init__(
+            self,
+            retriever=retriever,
+            reranker=reranker,
+            anthropic_api_key: str = "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",
+            model: str = "claude-3-5-sonnet-20241022",
+            max_tokens: int = 1024,
+            temperature: float = 0.0,
+    ):
+        self.retriever = retriever
+        self.reranker = reranker
+        self.client = Anthropic(api_key=anthropic_api_key)
+        self.model = model
+        self.max_tokens = max_tokens
+        self.temperature = temperature
+
+    def get_context(self, query: str) -> List[str]:
+        # Get initial candidates from retriever
+        retrieved_docs = self.retriever.search_similar(query)
+
+        # Rerank the candidates
+        context = self.reranker.rerank(query, retrieved_docs)
+
+        return context
+
+    def generate_prompt(self, context: List[str]) -> str:
+        context = "\n".join(context)
+        prompt = f"""
+                    "אתה רופא שיניים, דובר עברית בלבד. קוראים לך 'רופא השיניים האלקטרוני העברי הראשון', ענה למטופל על השאלה שלו על סמך הקונטקס הבא: {context}. הוסף כמה שיותר פרטים, ודאג שהתחביר יהיה תקין ויפה. תעצור כשאתה מרגיש שמיצית את עצמך. אל תמציא דברים. ואל תענה בשפות שהן לא עברית.
+                 """
+        return prompt
+
+    def get_response(self, question: str) -> str:
+        # Get relevant context
+        context = self.get_context(question)
+
+        # Generate prompt with context
+        prompt = self.generate_prompt(context)
+
+        # Get response from Claude
+        response = self.client.messages.create(
+            model=self.model,
+            max_tokens=self.max_tokens,
+            temperature=self.temperature,
+            messages=[
+                        {"role": "assistant", "content": prompt},
+                        {"role": "user", "content": f"{question}"}
+                    ]
+        )
+
+        return response.content[0].text

reranker.py

@@ -0,0 +1,22 @@
+from sklearn.metrics.pairwise import cosine_similarity
+
+
+from text_embedder_encoder import TextEmbedder
+
+
+class Reranker:
+    def __init__(self):
+        self.text_embedder = TextEmbedder()
+
+    def rerank(self, query, retrieved_docs, top_n=5):
+        # Encode query and documents
+        query_embedding = self.text_embedder.encode(query)
+        doc_embeddings = self.text_embedder.encode_many(retrieved_docs)
+        similarity_scores = cosine_similarity([query_embedding], doc_embeddings)[0]
+
+        similarity_scores_with_idxes = list(zip(similarity_scores, range(len(similarity_scores))))
+        similarity_scores_with_idxes.sort(reverse=True)
+        similarity_scores_with_idxes_final = similarity_scores_with_idxes[:top_n]
+        reranked_docs = [retrieved_docs[idx] for score, idx in similarity_scores_with_idxes_final if score >= 0.7]
+
+        return reranked_docs

retriever.py

@@ -0,0 +1,43 @@
+from pinecone import Pinecone
+
+
+from text_embedder_encoder import TextEmbedder, encoder_model_name
+
+
+class Retriever:
+    def __init__(self,
+                 pinecone_api_key="pcsk_468XZz_QfKbP3dWCh6nLatJjd882DGF5HDh6TupzEAeRpFLAMtDfDiPDNRC537Q4jAtxhV",
+                 index_name=f"hebrew-dentist-qa-{encoder_model_name.replace('/', '-')}".lower()):
+        # Initialize Pinecone connection
+        self.pc = Pinecone(api_key=pinecone_api_key)
+        self.index_name = index_name
+        self.text_embedder = TextEmbedder()
+        self.vector_dim = 768
+
+    def search_similar(self, query_text, top_k=50):
+        """
+        Search for similar content using vector similarity in Pinecone
+        """
+        try:
+            # Generate embedding for query
+            query_vector = self.text_embedder.encode(query_text)
+
+            # Get Pinecone index
+            index = self.pc.Index(self.index_name)
+
+            # Execute search
+            results = index.query(
+                vector=query_vector,
+                top_k=top_k,
+                include_metadata=True,
+            )
+
+            answers = []
+            for match in results['matches']:
+                answer = match['metadata']['answer']
+                answers.append(answer)
+
+            return answers
+        except Exception as e:
+            print(f"Error performing similarity search: {e}")
+            return []

text_embedder_encoder.py

@@ -0,0 +1,56 @@
+import torch
+import numpy as np
+from sentence_transformers import SentenceTransformer
+from typing import List
+
+
+encoder_model_name = 'MPA/sambert'
+
+
+class TextEmbedder:
+    def __init__(self):
+        """
+        Initialize the Hebrew text embedder using dictabert-large-heq model
+        """
+        # self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = SentenceTransformer(encoder_model_name)
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model.to(self.device)
+        self.model.eval()
+
+    def encode(self, text) -> np.ndarray:
+        """
+            Encode Hebrew text using LaBSE model with handling for texts longer than max_seq_length.
+
+            Args:
+                text (str): Hebrew text to encode
+                model_name (str): Name of the model to use
+                # max_seq_length (int): Maximum sequence length for the model
+                strategy (str): Strategy for combining sentence embeddings ('mean' or 'concat')
+
+            Returns:
+                numpy.ndarray: Text embedding
+            """
+        # Get embeddings for the text
+        embeddings = [float(x) for x in self.model.encode([text])[0]]
+
+        return embeddings
+
+    def encode_many(self, texts: List[str]) -> np.ndarray:
+        """
+            Encode Hebrew text using LaBSE model with handling for texts longer than max_seq_length.
+
+            Args:
+                text (str): Hebrew text to encode
+                model_name (str): Name of the model to use
+                # max_seq_length (int): Maximum sequence length for the model
+                strategy (str): Strategy for combining sentence embeddings ('mean' or 'concat')
+
+            Returns:
+                numpy.ndarray: Text embedding
+            """
+        # Get embeddings for the text
+        embeddings = self.model.encode(texts)
+        embeddings = [[float(x) for x in embedding] for embedding in embeddings]
+
+        return embeddings