Update app.py

app.py

@@ -1,25 +1,19 @@
 import gradio as gr
 import torch
 from sentence_transformers import SentenceTransformer, util
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import PyPDF2
 import os
 import time
 import logging
+from yacana import Agent, Task, LoggerManager
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
+yacana_logger = LoggerManager()
 
-# Load models with error handling
-try:
-    retriever_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
-    gen_tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-1")
-    gen_model = AutoModelForCausalLM.from_pretrained("microsoft/phi-1", torch_dtype=torch.float16)
-    logger.info("Models loaded successfully")
-except Exception as e:
-    logger.error(f"Error loading models: {str(e)}")
-    raise
+# Load retriever model
+retriever_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
 
 # Cache for document embeddings
 embedding_cache = {}
@@ -57,7 +51,7 @@ def chunk_text(text, chunk_size=500):
     return chunks
 
 def get_document_embeddings(documents):
-    """Compute embeddings for documents, using cache if available, and return a stacked tensor."""
+    """Compute embeddings for documents, using cache if available."""
     embeddings = []
     for doc in documents:
         if doc in embedding_cache:
@@ -68,25 +62,23 @@ def get_document_embeddings(documents):
             embeddings.append(emb)
     return torch.stack(embeddings)
 
-def generate_response(prompt):
-    """Helper function to generate text with Phi-1."""
-    try:
-        logger.info(f"Generating response for prompt: {prompt[:100]}...")  # Log first 100 chars
-        inputs = gen_tokenizer(prompt, return_tensors="pt")
-        logger.info(f"Tokenized input shape: {inputs['input_ids'].shape}")
-        outputs = gen_model.generate(**inputs, max_new_tokens=150, num_beams=2)
-        logger.info(f"Generated output shape: {outputs.shape}")
-        return gen_tokenizer.decode(outputs[0], skip_special_tokens=True)
-    except Exception as e:
-        logger.error(f"Error in generate_response: {str(e)}")
-        return f"Generation error: {str(e)}"
+def retrieve_context(question, documents):
+    """Retrieve top 3 relevant chunks."""
+    doc_embeddings = get_document_embeddings(documents)
+    query_embedding = retriever_model.encode(question, convert_to_tensor=True)
+    cos_scores = util.pytorch_cos_sim(query_embedding, doc_embeddings)[0]
+    top_results = torch.topk(cos_scores, k=min(3, len(documents)))
+    retrieved_context = ""
+    for score, idx in zip(top_results.values, top_results.indices):
+        retrieved_context += f"- {documents[idx]} (score: {score:.2f})\n"
+    return retrieved_context
 
 def rag_pipeline(question, pdf_files):
-    """RAG pipeline with multi-step thinking using Phi-1."""
+    """RAG pipeline with Yacana and Phi-1."""
     start_time = time.time()
     documents = []
 
-    # Process PDFs if provided
+    # Process PDFs or use default documents
     if pdf_files:
         for pdf in pdf_files:
             pages = extract_text_from_pdf(pdf)
@@ -94,7 +86,6 @@ def rag_pipeline(question, pdf_files):
                 chunks = chunk_text(page)
                 documents.extend(chunks)
     else:
-        # Default documents relevant to AI and Data Science
         documents = [
             "Artificial Intelligence (AI) is the simulation of human intelligence in machines.",
             "Data Science involves extracting insights from structured and unstructured data using statistical methods.",
@@ -105,55 +96,57 @@ def rag_pipeline(question, pdf_files):
     if not documents:
         return "No valid text could be extracted from the PDFs."
 
-    # Compute embeddings with caching
-    doc_embeddings = get_document_embeddings(documents)
-
-    # Embed the query
-    query_embedding = retriever_model.encode(question, convert_to_tensor=True)
-
-    # Retrieve top 3 chunks using cosine similarity
-    cos_scores = util.pytorch_cos_sim(query_embedding, doc_embeddings)[0]
-    top_results = torch.topk(cos_scores, k=min(3, len(documents)))
-    retrieved_context = ""
-    for score, idx in zip(top_results.values, top_results.indices):
-        retrieved_context += f"- {documents[idx]} (score: {score:.2f})\n"
-    
+    # Retrieve context
+    retrieved_context = retrieve_context(question, documents)
     logger.info(f"Retrieved context:\n{retrieved_context}")
 
-    # Step 1: Initial Answer
-    initial_prompt = (
-        f"Using the following context, provide a concise answer to the question:\n\n"
-        f"Context:\n{retrieved_context}\n\n"
-        f"Question: {question}\n\n"
-        f"Answer:"
+    # Define Yacana agents and tasks
+    agent = Agent("Phi1Agent", "phi", logger=yacana_logger)  # Assumes phi-1 via Ollama
+
+    # Task 1: Initial Answer
+    initial_task = Task(
+        name="GenerateInitialAnswer",
+        instruction=(
+            f"Using the following context, provide a concise answer to the question:\n\n"
+            f"Context:\n{retrieved_context}\n\n"
+            f"Question: {question}\n\n"
+            f"Answer:"
+        ),
+        agent=agent
     )
-    initial_answer = generate_response(initial_prompt)
-    if "Generation error" in initial_answer:
-        return initial_answer
-
-    # Step 2: Refine Answer
-    refine_prompt = (
-        f"Given the context and initial answer, refine and improve the response to the question:\n\n"
-        f"Context:\n{retrieved_context}\n\n"
-        f"Question: {question}\n\n"
-        f"Initial Answer: {initial_answer}\n\n"
-        f"Refined Answer:"
+
+    # Task 2: Refine Answer
+    initial_result = initial_task.run()
+    refine_task = Task(
+        name="RefineAnswer",
+        instruction=(
+            f"Given the context and initial answer, refine and improve the response:\n\n"
+            f"Context:\n{retrieved_context}\n\n"
+            f"Question: {question}\n\n"
+            f"Initial Answer: {initial_result}\n\n"
+            f"Refined Answer:"
+        ),
+        agent=agent
     )
-    refined_answer = generate_response(refine_prompt)
-    if "Generation error" in refined_answer:
-        return refined_answer
 
-    logger.info(f"Initial answer: {initial_answer}")
-    logger.info(f"Refined answer: {refined_answer}")
+    # Execute tasks
+    try:
+        refined_result = refine_task.run()
+        logger.info(f"Initial answer: {initial_result}")
+        logger.info(f"Refined answer: {refined_result}")
+    except Exception as e:
+        logger.error(f"Error in Yacana tasks: {str(e)}")
+        return f"Task execution error: {str(e)}"
+
     logger.info(f"Processing time: {time.time() - start_time:.2f} seconds")
-    return refined_answer if refined_answer else "Unable to generate a meaningful response."
+    return refined_result if refined_result else "Unable to generate a meaningful response."
 
 # Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("# RAG Pipeline with microsoft/phi-1 and Multi-Step Thinking")
+    gr.Markdown("# RAG Pipeline with microsoft/phi-1 and Yacana")
     gr.Markdown(
         "Upload PDFs (or use default AI/Data Science docs), ask a question, "
-        "and get refined answers using Phi-1 with multi-step reasoning on 2 vCPUs and 16GB RAM."
+        "and get refined answers using Phi-1 via Yacana on 2 vCPUs and 16GB RAM."
     )
     with gr.Row():
         with gr.Column():