RAG Placeholder demo test

app.py

@@ -76,9 +76,120 @@ from sentence_transformers import SentenceTransformer
 #         yield response, f"{tokens_per_second:.2f}"
 
 
-# Initialize GPU tensor
-zero = torch.Tensor([0]).cuda()
-print(zero.device)  # This will print 'cpu' outside the @spaces.GPU decorated function
+#---------
+#----------
+
+# # Initialize GPU tensor
+# zero = torch.Tensor([0]).cuda()
+# print(zero.device)  # This will print 'cpu' outside the @spaces.GPU decorated function
+
+# # Load the embedding model
+# embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
+
+# # Load the Qwen model and tokenizer
+# llmguide_model = AutoModelForCausalLM.from_pretrained(
+#     "Qwen/Qwen2-0.5B-Instruct",
+#     torch_dtype="auto",
+#     device_map="auto"
+# )
+# llmguide_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
+
+# # Sample knowledge base (replace with your own data)
+# knowledge_base = [
+#     "The capital of France is Paris.",
+#     "Python is a popular programming language.",
+#     "Machine learning is a subset of artificial intelligence.",
+#     "The Earth orbits around the Sun.",
+#     "orbits are a group of fans of a music group"
+# ]
+
+# # Create embeddings for the knowledge base
+# knowledge_base_embeddings = embedding_model.encode(knowledge_base)
+
+# def retrieve(query, k=2):
+#     query_embedding = embedding_model.encode([query])
+#     similarities = torch.nn.functional.cosine_similarity(torch.tensor(query_embedding), torch.tensor(knowledge_base_embeddings))
+#     top_k_indices = similarities.argsort(descending=True)[:k]
+#     return [knowledge_base[i] for i in top_k_indices]
+
+# def get_resource_usage():
+#     ram_usage = psutil.virtual_memory().percent
+#     gpu_memory_allocated = torch.cuda.memory_allocated() / (1024 ** 3)  # Convert to GB
+#     gpu_memory_reserved = torch.cuda.memory_reserved() / (1024 ** 3)  # Convert to GB
+#     return f"RAM Usage: {ram_usage:.2f}%, GPU Memory Allocated: {gpu_memory_allocated:.2f}GB, GPU Memory Reserved: {gpu_memory_reserved:.2f}GB"
+
+# @spaces.GPU
+# def llmguide_generate_response(prompt, stream=False):
+#     print(zero.device)  # This will print 'cuda:0' inside the @spaces.GPU decorated function
+    
+#     messages = [
+#         {"role": "system", "content": "You are a helpful assistant."},
+#         {"role": "user", "content": prompt}
+#     ]
+#     text = llmguide_tokenizer.apply_chat_template(
+#         messages,
+#         tokenize=False,
+#         add_generation_prompt=True
+#     )
+#     model_inputs = llmguide_tokenizer([text], return_tensors="pt").to(zero.device)
+
+#     start_time = time.time()
+#     total_tokens = 0
+
+#     if stream:
+#         streamer = TextIteratorStreamer(llmguide_tokenizer, skip_special_tokens=True)
+#         generation_kwargs = dict(
+#             model_inputs,
+#             streamer=streamer,
+#             max_new_tokens=512,
+#             temperature=0.7,
+#         )
+#         thread = Thread(target=llmguide_model.generate, kwargs=generation_kwargs)
+#         thread.start()
+
+#         generated_text = ""
+#         for new_text in streamer:
+#             generated_text += new_text
+#             total_tokens += 1
+#             current_time = time.time()
+#             tokens_per_second = total_tokens / (current_time - start_time)
+#             yield generated_text, f"{tokens_per_second:.2f}", ""
+
+#         resource_usage = get_resource_usage()
+#         yield generated_text, f"{tokens_per_second:.2f}", resource_usage
+#     else:
+#         generated_ids = llmguide_model.generate(
+#             model_inputs.input_ids,
+#             max_new_tokens=512
+#         )
+#         generated_ids = [
+#             output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+#         ]
+#         response = llmguide_tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+#         total_tokens = len(generated_ids[0])
+#         end_time = time.time()
+#         tokens_per_second = total_tokens / (end_time - start_time)
+#         resource_usage = get_resource_usage()
+#         yield response, f"{tokens_per_second:.2f}", resource_usage
+
+# def rag(query, stream=False):
+#     retrieved_docs = retrieve(query)
+#     context = " ".join(retrieved_docs)
+#     prompt = f"Context: {context}\nQuestion: {query}\nAnswer:"
+    
+#     generator = llmguide_generate_response(prompt, stream)
+    
+#     if stream:
+#         def stream_output():
+#             for generated_text, tokens_per_second, ram_usage in generator:
+#                 yield generated_text, tokens_per_second, ram_usage
+#         return stream_output()
+#     else:
+#         # For non-streaming, we just need to get the final output
+#         for generated_text, tokens_per_second, ram_usage in generator:
+#             pass  # This will iterate to the last yield
+#         return generated_text, tokens_per_second, ram_usage
+
 
 # Load the embedding model
 embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
@@ -93,31 +204,28 @@ llmguide_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
 
 # Sample knowledge base (replace with your own data)
 knowledge_base = [
-    "The capital of France is Paris.",
-    "Python is a popular programming language.",
-    "Machine learning is a subset of artificial intelligence.",
-    "The Earth orbits around the Sun.",
+    {"id": "doc1", "content": "The capital of France is Paris."},
+    {"id": "doc2", "content": "Python is a popular programming language."},
+    {"id": "doc3", "content": "Machine learning is a subset of artificial intelligence."},
+    {"id": "doc4", "content": "The Earth orbits around the Sun."},
+    {"id": "doc5", "content": "orbits is the name of a korean fangroup"},
 ]
 
 # Create embeddings for the knowledge base
-knowledge_base_embeddings = embedding_model.encode(knowledge_base)
+knowledge_base_embeddings = embedding_model.encode([doc["content"] for doc in knowledge_base])
 
 def retrieve(query, k=2):
     query_embedding = embedding_model.encode([query])
     similarities = torch.nn.functional.cosine_similarity(torch.tensor(query_embedding), torch.tensor(knowledge_base_embeddings))
     top_k_indices = similarities.argsort(descending=True)[:k]
-    return [knowledge_base[i] for i in top_k_indices]
+    return [(knowledge_base[i]["content"], knowledge_base[i]["id"]) for i in top_k_indices]
 
-def get_resource_usage():
-    ram_usage = psutil.virtual_memory().percent
-    gpu_memory_allocated = torch.cuda.memory_allocated() / (1024 ** 3)  # Convert to GB
-    gpu_memory_reserved = torch.cuda.memory_reserved() / (1024 ** 3)  # Convert to GB
-    return f"RAM Usage: {ram_usage:.2f}%, GPU Memory Allocated: {gpu_memory_allocated:.2f}GB, GPU Memory Reserved: {gpu_memory_reserved:.2f}GB"
+def get_ram_usage():
+    ram = psutil.virtual_memory()
+    return f"RAM Usage: {ram.percent:.2f}%, Available: {ram.available / (1024 ** 3):.2f}GB, Total: {ram.total / (1024 ** 3):.2f}GB"
 
 @spaces.GPU
-def llmguide_generate_response(prompt, stream=False):
-    print(zero.device)  # This will print 'cuda:0' inside the @spaces.GPU decorated function
-    
+def llmguide_generate_response(prompt, doc_ids=None, stream=False):
     messages = [
         {"role": "system", "content": "You are a helpful assistant."},
         {"role": "user", "content": prompt}
@@ -127,7 +235,7 @@ def llmguide_generate_response(prompt, stream=False):
         tokenize=False,
         add_generation_prompt=True
     )
-    model_inputs = llmguide_tokenizer([text], return_tensors="pt").to(zero.device)
+    model_inputs = llmguide_tokenizer([text], return_tensors="pt").to(llmguide_model.device)
 
     start_time = time.time()
     total_tokens = 0
@@ -149,10 +257,10 @@ def llmguide_generate_response(prompt, stream=False):
             total_tokens += 1
             current_time = time.time()
             tokens_per_second = total_tokens / (current_time - start_time)
-            yield generated_text, f"{tokens_per_second:.2f}", ""
+            yield generated_text, f"{tokens_per_second:.2f}", "", ", ".join(doc_ids) if doc_ids else "N/A"
 
-        resource_usage = get_resource_usage()
-        yield generated_text, f"{tokens_per_second:.2f}", resource_usage
+        ram_usage = get_ram_usage()
+        yield generated_text, f"{tokens_per_second:.2f}", ram_usage, ", ".join(doc_ids) if doc_ids else "N/A"
     else:
         generated_ids = llmguide_model.generate(
             model_inputs.input_ids,
@@ -165,36 +273,32 @@ def llmguide_generate_response(prompt, stream=False):
         total_tokens = len(generated_ids[0])
         end_time = time.time()
         tokens_per_second = total_tokens / (end_time - start_time)
-        resource_usage = get_resource_usage()
-        yield response, f"{tokens_per_second:.2f}", resource_usage
-
-    # Clear CUDA cache
-    # torch.cuda.empty_cache()
-    # gc.collect()
-
-# def rag(query, stream=False):
-#     retrieved_docs = retrieve(query)
-#     context = " ".join(retrieved_docs)
-#     prompt = f"Context: {context}\nQuestion: {query}\nAnswer:"
-#     return llmguide_generate_response(prompt, stream)
-
-def rag(query, stream=False):
-    retrieved_docs = retrieve(query)
-    context = " ".join(retrieved_docs)
-    prompt = f"Context: {context}\nQuestion: {query}\nAnswer:"
+        ram_usage = get_ram_usage()
+        yield response, f"{tokens_per_second:.2f}", ram_usage, ", ".join(doc_ids) if doc_ids else "N/A"
+
+def process_query(query, use_rag, stream=False):
+    if use_rag:
+        retrieved_docs = retrieve(query)
+        context = " ".join([doc for doc, _ in retrieved_docs])
+        doc_ids = [doc_id for _, doc_id in retrieved_docs]
+        prompt = f"Context: {context}\nQuestion: {query}\nAnswer:"
+    else:
+        prompt = query
+        doc_ids = None
     
-    generator = llmguide_generate_response(prompt, stream)
+    generator = llmguide_generate_response(prompt, doc_ids, stream)
     
     if stream:
         def stream_output():
-            for generated_text, tokens_per_second, ram_usage in generator:
-                yield generated_text, tokens_per_second, ram_usage
+            for generated_text, tokens_per_second, ram_usage, doc_references in generator:
+                yield generated_text, tokens_per_second, ram_usage, doc_references
         return stream_output()
     else:
         # For non-streaming, we just need to get the final output
-        for generated_text, tokens_per_second, ram_usage in generator:
+        for generated_text, tokens_per_second, ram_usage, doc_references in generator:
             pass  # This will iterate to the last yield
-        return generated_text, tokens_per_second, ram_usage
+        return generated_text, tokens_per_second, ram_usage, doc_references
+
 
 #--------------------------------------------------------------------------------------------------------------------------------
 
@@ -838,24 +942,43 @@ with gr.Blocks() as demo:
         <div style="width: 20%; text-align: center">HF + Gradio allows for api use so this my prototype tool for tool use test</div>
     </div>""")
     with gr.Accordion("Qwen 0.5B as Space Guide Tests", open=False):
-        gr.HTML("Placeholder for FAQ type - front end as prompt engineering for the first message to force direction of conversion")
         gr.HTML("Placeholder for weak RAG Type - Explanations through an opensource embeddings engine")
+        # gr.Interface(
+        #     fn=rag,
+        #     inputs=[
+        #         gr.Textbox(lines=2, placeholder="Enter your question here..."),
+        #         gr.Checkbox(label="Stream output")
+        #     ],
+        #     outputs=[
+        #         gr.Textbox(label="Generated Response"),
+        #         gr.Textbox(label="Tokens per second"),
+        #         gr.Textbox(label="Resource Usage")
+        #     ],
+        #     title="RAG Q&A System with GPU Acceleration and Resource Monitoring",
+        #     description="Ask a question and get an answer based on the retrieved context. The response is generated using a GPU-accelerated model. Resource usage is logged at the end of generation."
+        # )
+
         gr.Interface(
-            fn=rag,
+            fn=process_query,
             inputs=[
                 gr.Textbox(lines=2, placeholder="Enter your question here..."),
+                gr.Checkbox(label="Use RAG"),
                 gr.Checkbox(label="Stream output")
             ],
             outputs=[
                 gr.Textbox(label="Generated Response"),
                 gr.Textbox(label="Tokens per second"),
-                gr.Textbox(label="Resource Usage")
+                gr.Textbox(label="RAM Usage"),
+                gr.Textbox(label="Referenced Documents")
             ],
-            title="RAG Q&A System with GPU Acceleration and Resource Monitoring",
-            description="Ask a question and get an answer based on the retrieved context. The response is generated using a GPU-accelerated model. Resource usage is logged at the end of generation."
+            title="RAG/Non-RAG Q&A System",
+            description="Ask a question with or without using RAG. The response is generated using a GPU-accelerated model. RAM usage and referenced document IDs (for RAG) are logged."
         )
-        ("Placeholder for https://huggingface.co/h2oai/h2o-danube3-500m-chat-GGUF as alternative")
-        ("Placeholder for qwen 2 72b as alternative use checkbox and gradio client api call")
+        gr.HTML("Placeholder for FAQ type (merge as buttons on the above interface) - front end as prompt engineering for the first message to force direction of conversion")
+        
+
+        gr.HTML("Placeholder for https://huggingface.co/h2oai/h2o-danube3-500m-chat-GGUF as alternative")
+        gr.HTML("Placeholder for qwen 2 72b as alternative use checkbox and gradio client api call")
         gr.Markdown("# Qwen-0.5B-Instruct Language Model")
         gr.Markdown("This demo uses the Qwen-0.5B-Instruct model to generate responses based on your input.")
         gr.HTML("Example prompts: <br>I am writing a story about a chef. please write dishes to appear on the menu. <br>What are the most common decisions that a chef story would include? <br>What are the kinds problems that a chef story would include? <br>What are the kinds of out of reach goals that a chef story would include? <br>Continue this config - Paste any complete block of the config")