Update app.py

app.py

@@ -1,65 +1,177 @@
-import gradio as gr
-from huggingface_hub import InferenceClient
-
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-#client = InferenceClient("Qwen/Qwen2.5-72B-Instruct")
-client = InferenceClient("Qwen/Qwen2.5-Coder-32B-Instruct")
-
-
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": 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": message})
-
-    response = ""
-
-    for message in 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
-
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = 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)",
-        ),
-    ],
-)
-
-
-if __name__ == "__main__":
-    demo.launch()
+import gradio as gr  
+import os  
+import sqlite3  
+import numpy as np  
+from huggingface_hub import InferenceClient  
+from sentence_transformers import SentenceTransformer, util  
+
+
+# Initialize the model and InferenceClient  
+model = SentenceTransformer("xmanii/maux-gte-persian", trust_remote_code=True)  
+
+huggingface_token = os.getenv('RAG') 
+print(huggingface_token)
+
+client = InferenceClient(api_key = huggingface_token)
+
+
+# Database initialization  
+DB_NAME = 'files.db'  
+conn = sqlite3.connect(DB_NAME, check_same_thread=False)  
+cursor = conn.cursor()  
+
+# Creating a table to store file information if it doesn't exist  
+cursor.execute('''  
+    CREATE TABLE IF NOT EXISTS files (  
+        id INTEGER PRIMARY KEY AUTOINCREMENT,  
+        filename TEXT NOT NULL,  
+        path TEXT NOT NULL,  
+        chunks TEXT NOT NULL  
+    )  
+''')  
+
+# Creating a table to store embeddings  
+cursor.execute('''  
+    CREATE TABLE IF NOT EXISTS embeddings (  
+        id INTEGER PRIMARY KEY AUTOINCREMENT,  
+        file_id INTEGER NOT NULL,  
+        chunk_index INTEGER NOT NULL,  
+        embedding BLOB NOT NULL,  
+        FOREIGN KEY (file_id) REFERENCES files(id)  
+    )  
+''')  
+conn.commit()  
+
+UPLOAD_FOLDER = 'uploads'  
+os.makedirs(UPLOAD_FOLDER, exist_ok=True)  
+
+def read_file(file_path):  
+    try:  
+        with open(file_path, 'r', encoding='utf-8') as file:  
+            return file.read()  
+    except FileNotFoundError:  
+        return None  
+    except Exception as e:  
+        return None  
+
+def chunk_text(text, chunk_size=1000, overlap_size=100):  
+    chunks = []  
+    start = 0  
+    while start < len(text):  
+        end = start + chunk_size  
+        chunks.append(text[start:end])  
+        start += (chunk_size - overlap_size)  
+    return chunks  
+
+def calculate_similarity(embedding1, embedding2):  
+    similarity = util.cos_sim(embedding1, embedding2)  
+    return similarity.item()  
+
+def store_embedding(file_id, chunk_index, embedding):  
+    cursor.execute("INSERT INTO embeddings (file_id, chunk_index, embedding) VALUES (?, ?, ?)",  
+                   (file_id, chunk_index, embedding))  
+    conn.commit()  
+
+def get_embeddings(file_id):  
+    cursor.execute("SELECT chunk_index, embedding FROM embeddings WHERE file_id = ?", (file_id,))  
+    return cursor.fetchall()  
+
+def upload_file(file):  
+    if file is None:  
+        return "No file selected.", None  
+
+    file_path = os.path.join(UPLOAD_FOLDER, file.name)  
+    with open(file_path, 'wb') as f:  
+        f.write(file.read())  
+
+    text = read_file(file_path)  
+    if text is None:  
+        return f"Could not read file at path: {file_path}", None  
+
+    chunks = chunk_text(text)  
+    chunks_str = '|'.join(chunks)  
+
+    # Insert file metadata  
+    cursor.execute("INSERT INTO files (filename, path, chunks) VALUES (?, ?, ?)",  
+                   (file.name, file_path, chunks_str))  
+    file_id = cursor.lastrowid  # get the id of the inserted file  
+    conn.commit()  
+
+    print("Calculate and store embeddings")  
+    for index, chunk in enumerate(chunks):  
+        embedding = model.encode(chunk).tobytes()  # convert to bytes for storage  
+        store_embedding(file_id, index, embedding)  
+
+    return f"File '{file.name}' uploaded and processed successfully.", file_id  
+
+def chat(input_sentence, file_id):  
+    if not input_sentence:  
+        return "User input is required."  
+
+    # Get the embeddings from the uploaded file  
+    cursor.execute("SELECT id FROM files WHERE id = ?", (file_id,))  
+    if cursor.fetchone() is None:  
+        return "No uploaded files found. Please upload a file."  
+
+    embeddings = get_embeddings(file_id)  
+    if not embeddings:  
+        return "No embeddings found for the uploaded file."  
+
+    # Encode the user input for comparison  
+    input_embedding = model.encode(input_sentence)  
+
+    # Calculate similarities using stored embeddings  
+    similarities = []  
+    for index, embedding in embeddings:  
+        stored_embedding = np.frombuffer(embedding, dtype=np.float32)  # convert bytes back to numpy array  
+        similarity = calculate_similarity(input_embedding, stored_embedding)  
+        similarities.append(similarity)  
+
+    sorted_indices = sorted(range(len(similarities)), key=lambda i: similarities[i], reverse=True)  
+
+    # Get the top chunks based on similarity  
+    cursor.execute("SELECT chunks FROM files WHERE id = ?", (file_id,))  
+    chunks = cursor.fetchone()[0].split('|')  
+    answercontext = "".join(chunks[i] for i in sorted_indices[:5])  
+
+    system_prompt = (  
+        "فقط بر اساس متن زیر به پرسش کاربر پاسخ بده و اگر پاسخ در متن نبود بنویس نمیدانم.\n\n"  
+        "متن:\n" + answercontext)  
+    
+    messages = [  
+        {"role": "system", "content": system_prompt},  
+        {"role": "user", "content": input_sentence}  
+    ]  
+
+    response = ""  
+    for message in client.chat_completion(  
+        messages,  
+        temperature=0.5,  
+        max_tokens=4096,  
+        top_p=0.7,  
+        stream=False  
+    ):  
+        response += message.choices[0].delta.content  
+
+    return response  
+
+# Gradio Interface  
+with gr.Blocks() as demo:  
+    gr.Markdown("## File Upload and Chat Interface")  
+    
+    with gr.Row():  
+        file_input = gr.File(label="Upload a file")  
+        upload_button = gr.Button("Upload File")  
+    
+    upload_output = gr.Textbox(label="Upload Status", interactive=False)  
+    file_id = gr.State()  
+
+    upload_button.click(upload_file, inputs=file_input, outputs=[upload_output, file_id])  
+
+    user_input = gr.Textbox(label="Your Question")  
+    chat_button = gr.Button("Ask")  
+    chat_output = gr.Textbox(label="Response", interactive=False)  
+
+    chat_button.click(chat, inputs=[user_input, file_id], outputs=chat_output)  
+
+demo.launch()