update

app.py

@@ -0,0 +1,69 @@
+import os
+import gradio as gr
+from transformers import pipeline
+from elasticsearch import Elasticsearch
+
+# Connect to Elasticsearch
+es = Elasticsearch(hosts=["https://data.neuralnoise.com:9200"],
+                   basic_auth=('elastic', os.environ['ES_PASSWORD']),
+                   verify_certs=False, ssl_show_warn=False)
+
+# Load your language model from HuggingFace Transformers
+generator = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.2")
+
+def search_es(query, index="pubmed", num_results=3):
+    """
+    Search the Elasticsearch index for the most relevant documents.
+    """
+    
+    print(f'Running query: {query}')
+    
+    response = es.search(
+        index=index,
+        body={
+            "query": {
+                "match": {
+                    "content": query  # Assuming documents have a 'content' field
+                }
+            },
+            "size": num_results
+        }
+    )
+    
+    # Extract and return the documents
+    docs = [hit["_source"]["content"] for hit in response['hits']['hits']]
+
+    print(f'Received {len(docs)} documents')
+
+    return docs
+
+def rag_pipeline(prompt, index="pubmed"):
+    """
+    A simple RAG pipeline that retrieves documents and uses them to enrich the context for the LLM.
+    """
+    # Retrieve documents
+    docs = search_es(prompt, index=index)
+    
+    # Combine prompt with retrieved documents
+    enriched_prompt = f"{prompt}\n\n{' '.join(docs)}"
+    
+    # Generate response using the LLM
+    response = generator(enriched_prompt, max_new_tokens=256, return_full_text=False)
+    
+    # Return the generated text and the documents
+    return response[0]['generated_text'], "\n\n".join(docs)
+
+# Create the Gradio interface
+iface = gr.Interface(fn=rag_pipeline,
+                     inputs=[
+                         gr.Textbox(label="Input Prompt"),
+                         gr.Textbox(label="Elasticsearch Index", value="pubmed")  # Corrected here
+                     ],
+                     outputs=[
+                         gr.Textbox(label="Generated Text"),
+                         gr.Textbox(label="Retrieved Documents")
+                     ],
+                     description="Retrieval-Augmented Generation Pipeline")
+
+# Launch the interface
+iface.launch()

requirements.txt

@@ -0,0 +1,3 @@
+torch
+transformers
+elasticsearch