Update app.py

app.py

@@ -5,42 +5,15 @@ import chromadb
 import torch
 from sentence_transformers import SentenceTransformer
 import os
+from chromadb.utils import embedding_functions
 
-# ChromaDB Setup (Persistent Client)
-CHROMA_DB_PATH = "new_hadith_rag_source"
-client = chromadb.PersistentClient(path=CHROMA_DB_PATH)
-COLLECTION_NAME = "hadiths_new_complete"
-
-def load_or_create_collection():
-    try:
-        collection = client.get_collection(name=COLLECTION_NAME)
-        print("Collection loaded successfully.")
-        return collection
-    except ValueError:
-        print("Creating new collection...")
-        collection = client.create_collection(name=COLLECTION_NAME, overwrite=True)
-        ds = load_dataset("rwmasood/hadith-qa-pair")
-        device = 'cpu'
-        embedding_model = SentenceTransformer('all-MiniLM-L6-v2').to(device)
-        
-        for split in ds.keys():
-            documents = [
-                f"Hadith: {row['hadith-eng']}\nQuestion: {row['question']}\nReference: {row['reference']}"
-                for row in ds[split]
-            ]
-            ids = [f"{split}_{i}" for i in range(len(documents))]
-            embeddings = embedding_model.encode(documents, convert_to_tensor=True, device=device).numpy()
-            
-            collection.add(
-                documents=documents,
-                ids=ids,
-                embeddings=embeddings
-            )
-        
-        print(f"Collection created with {collection.count()} documents.")
-        return collection
-
-collection = load_or_create_collection()
+# Initialize ChromaDB client with the existing path
+client = chromadb.PersistentClient(path="new_hadith_rag_source")
+
+# Load the existing collection
+collection = client.get_collection(name="hadiths_new_complete")
+
+# Debugging print to verify the number of documents in the collection
 print(f"Number of documents in collection: {collection.count()}")
 
 # Model and Tokenizer Loading
@@ -54,74 +27,99 @@ llm = AutoModelForSeq2SeqLM.from_pretrained(
     device_map="auto"
 )
 
-# Embedding Model
-device = 'cpu'
-model = SentenceTransformer('all-MiniLM-L6-v2').to(device)
+# Load the pre-trained model and tokenizer
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+retrieval_model = SentenceTransformer('all-MiniLM-L6-v2').to(device)
 
-def query_collection(query_text, top_k=3):
-    query_embedding = model.encode(query_text, convert_to_tensor=True, device=device).numpy()
-    results = collection.query(
-        query_embeddings=[query_embedding],
-        n_results=top_k
-    )
+# Function to query the collection
+def query_collection(query, n_results):
+    # Compute the embedding for the query
+    query_embedding = retrieval_model.encode([query], convert_to_tensor=True, device=device).cpu().numpy()
+    
+    # Query the collection
+    results = collection.query(query_embeddings=query_embedding, n_results=n_results)
+    
     return results
 
-def speculative_generation(context, question, num_candidates=3):
-    responses = []
-    for _ in range(num_candidates):
-        prompt = f"Context:\n{context}\n\nQuestion:\n{question}\n\nAnswer:"
-        inputs = tokenizer(prompt, return_tensors="pt").to(device)
-        try:
-            outputs = llm.generate(**inputs, max_length=2048, num_return_sequences=1, num_beams=5, temperature=0.9, pad_token_id=tokenizer.eos_token_id)
-            response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-            responses.append(response)
-        except Exception as e:
-            print(f"Error during generation: {e}")
-            responses.append(f"An error occurred during generation: {e}")
-    return responses
-
-def grade_responses(responses, query):
-    best_score = -1
-    best_response = ""
-    for response in responses:
-        score = sum(1 for word in query.lower().split() if word in response.lower())
-        query_embedding = model.encode(query, convert_to_tensor=True, device=device)
-        response_embedding = model.encode(response, convert_to_tensor=True, device=device)
-        similarity = torch.nn.functional.cosine_similarity(query_embedding, response_embedding, dim=0).item()
-        score += similarity * 10
-        if score > best_score:
-            best_score = score
-            best_response = response
-    return best_response
-
-def chatbot_response(user_query, top_k=3, num_candidates=3):
+# Generate a response using the retrieved documents as context
+def generate_response(context, question):
+    prompt = f"Please provide a short, well-structured answer and avoids repetition from context:\n{context}\n\nQuestion:\n{question}\n\nAnswer:"
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    outputs = llm.generate(**inputs, max_length=2048, num_return_sequences=1, num_beams=5, temperature=0.9, pad_token_id=tokenizer.eos_token_id)
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response
+
+# Main chatbot function with basic RAG
+def chatbot_response(user_query, top_k=2):
+    # Step 1: Retrieve relevant documents
     results = query_collection(user_query, top_k)
-    context = "\n\n".join(results['documents'][0])
-    speculative_responses = speculative_generation(context, user_query, num_candidates)
-    best_response = grade_responses(speculative_responses, user_query)
-    return best_response
+    
+    # Step 2: Combine retrieved documents into context
+    documents = [doc for doc_list in results['documents'] for doc in doc_list]
+    combined_context = "\n\n".join(documents)
+    
+    # Step 3: Generate a response using the combined context
+    response = generate_response(combined_context, user_query)
+    
+    return response
+
+# Global variable to control the processing state
+stop_processing = False
 
 def chatbot(query, num_candidates):
+    global stop_processing
+    stop_processing = False  # Reset stop flag at the beginning of each query
+
+    # Jika query kosong, kembalikan pesan default
     if not query.strip():
         return "Please ask a question about hadiths."
-    try:
-        answer = chatbot_response(query, num_candidates)
-        if "don't know" in answer.lower() or "not sure" in answer.lower():
-            return "Sorry. I don't have information about the hadiths related. It might be a dhoif, or maudhu, or I just don't have the knowledge."
-        else:
-            return answer
-    except Exception as e:
-        print(f"Error in chatbot: {e}")
-        return f"An error occurred: {e}"
-
-if __name__ == "__main__":
-    iface = gr.Interface(
-        fn=chatbot,
-        inputs=[
-            gr.Textbox(lines=2, placeholder="Enter your question here..."),
-            gr.Slider(minimum=1, maximum=10, value=3, step=1, label="Number of Hadiths as References")],
-        outputs=gr.Textbox(label="Answer"),
-        title="Hadith QA Chatbot",
-        description="Ask questions related to Hadiths."
+    
+    # Lakukan retrieval dan generation dengan Speculative RAG
+    answer = chatbot_response(query, num_candidates)
+    
+    # Check if stop button was pressed
+    if stop_processing:
+        return "Processing was stopped by the user."
+
+    # Format jawaban
+    if "don't know" in answer.lower() or "not sure" in answer.lower():
+        return "Sorry. I don't have information about the hadiths related. It might be a dhoif, or maudhu, or I just don't have the knowledge."
+    else:
+        return answer
+
+def stop():
+    global stop_processing
+    stop_processing = True
+    return "Processing stopped."
+
+# Buat Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown(
+        """
+        # Burhan AI
+        Assalamualaikum! I am Burhan AI, a chatbot that can help you find answers to your questions about hadiths. 
+        \n 
+        Please note that this is a demo version and may not be perfect.
+        This chatbot is powered by the ChromaDB and Flan-T5-base models with RAG architecture.
+        Flan-T5-base is a small model and may not be as accurate as the bigger models.
+        If you have any feedback or suggestions, you can contact me at [email protected]
+        \n
+        Jazakallah Khairan!
+        """
     )
-    iface.launch()
+    with gr.Row():
+        query_input = gr.Textbox(lines=2, placeholder="Enter your question here...")
+        num_candidates_input = gr.Slider(minimum=1, maximum=10, value=2, step=1, label="Number of References")
+        submit_button = gr.Button("Submit")
+    
+    output_text = gr.Textbox(label="Response")
+
+    submit_button.click(chatbot, inputs=[query_input, num_candidates_input], outputs=output_text)
+
+    # Add a button to stop processing
+    stop_button = gr.Button("Stop Processing")
+    stop_output = gr.Textbox(visible=False)
+    stop_button.click(stop, inputs=[], outputs=stop_output)
+
+# Jalankan Gradio interface
+demo.launch()