updates

app.py

@@ -1,121 +1,61 @@
 from typing import Optional
-from sentence_transformers import SentenceTransformer
-import pymongo
-import os
-from huggingface_hub import InferenceClient
 import gradio as gr
+from build_rag import get_context
+from models.gemma import gemma_predict
+from models.gemini import get_gemini_response
 
 
-HF_token = os.environ["HF_TOKEN"]
+def clear():
+    return None, None, None
 
 
-def get_embedding(text: str) -> list[float]:
-    embedding_model = SentenceTransformer("thenlper/gte-large")
-
-    if not text.strip():
-        print("Attempted to get embedding for empty text.")
-        return []
-
-    embedding = embedding_model.encode(text)
-
-    return embedding.tolist()
-
-
-def get_mongo_client(mongo_url):
-    """Establish connection to the MongoDB."""
-    if not mongo_url:
-        print("MONGO_URI not set in environment variables")
-    try:
-        client = pymongo.MongoClient(mongo_url)
-        print("Connection to MongoDB successful")
-        return client
-    except pymongo.errors.ConnectionFailure as e:
-        print(f"Connection failed: {e}")
-        return None
-
-
-def get_mongo_url():
-    username = os.environ["MONGO_USERNAME"]
-    password = os.environ["MONGO_PW"]
-    mongo_url = f"mongodb+srv://{username}:{password}@cluster0.62unmco.mongodb.net/"
-    return mongo_url
-
-
-def query_results(query, mongo_url):
-    mongo_client = get_mongo_client(mongo_url)
-    db = mongo_client["EU_Cities"]
-
-    query_embedding = get_embedding(query)
-    results = db.EU_cities_collection.aggregate([
-        {
-            "$vectorSearch": {
-                "index": "vector_index",
-                "path": "embedding",
-                "queryVector": query_embedding,
-                "numCandidates": 150,
-                "limit": 5
-            }
-        }
-    ])
-    return results
-
-
-def get_search_result(query, mongo_url):
-    get_knowledge = query_results(query, mongo_url)
-    print(get_knowledge)
-
-    search_result = ""
-    for result in get_knowledge:
-        search_result += f"City: {result.get('city', 'N/A')}, Abstract: {result.get('combined', 'N/A')}\n"
-
-    return search_result
-
-
-def generate_text(query, model_name: Optional[str] = "google/gemma-2b-it"):
-    if model_name is None:
-        model_name = "google/gemma-2b-it"
-
-    mongo_url = get_mongo_url()
-    source_information = get_search_result(query, mongo_url)
-    combined_information = (
-        f"Query: {query}\nContinue to answer the query by using the Search Results:\n{source_information}."
-    )
-    client = InferenceClient(model_name, token=HF_token)
-
-    stream = client.text_generation(prompt=combined_information, details=True, stream=True, max_new_tokens=2048,
-                                    return_full_text=False)
-    output = ""
-
-    for response in stream:
-        output += response.token.text
-
-    if "<eos>" in output:
-        output = output.split("<eos>")[0]
-    return output
+def generate_text(query_text, model_name: Optional[str] = "google/gemma-2b-it"):
+    combined_information = get_context(query_text)
+    if model_name is None or model_name == "google/gemma-2b-it":
+        return gemma_predict(combined_information, model_name)
+    if model_name == "gemini-1.0-pro":
+        return get_gemini_response(combined_information, model_name, None)
+    return "Sorry, something went wrong! Please try again."
 
 
 examples = [["I'm planning a vacation to France. Can you suggest a one-week itinerary including must-visit places and "
-             "local cuisines to try?", None],
-            ["I want to explore off-the-beaten-path destinations in Europe, any suggestions?", None],
-            ["Suggest some cities that can be visited from London and are very rich in history and culture.", None],
+             "local cuisines to try?", "google/gemma-2b-it"],
+            # ["I want to explore off-the-beaten-path destinations in Europe, any suggestions?", "gemini-1.0-pro"],
+            ["Suggest some cities that can be visited from London and are very rich in history and culture.",
+             "google/gemma-2b-it"],
             ]
 
-demo = gr.Interface(
-    fn=generate_text,
-    inputs=["text",
-            gr.Dropdown(
-                ["google/gemma-2b-it"], label="Models", info="Will "
-                                                                                                             "add "
-                                                                                                             "more "
-                                                                                                             "models "
-                                                                                                             "later! "
-            ),
-            ],
-    title="🇪🇺 Euro TravelBot 🇪🇺",
-    description="Explore Europe with ease using our prototype app! We're testing the compatibility of RAG implementations with Google Gemma-2b-it models to generate travel recommendations. This early version (read quick and dirty implementation) aims to see if functionalities work smoothly. It relies on Wikipedia abstracts from 160 European cities to provide answers to your questions. Please be kind with it as it's a work in progress!",   
-    outputs=["text"],
-    examples=examples,
-)
+with gr.Blocks() as demo:
+    gr.HTML("""<center><h1 style='font-size:xx-large;'>🇪🇺 Euro City Recommender using Gemini & Gemma 🇪🇺</h1><br><h3>Gemini 
+    & Gemma Sprints 2024 submissions by Ashmi Banerjee. </h3></center> <br><p>We're testing the compatibility of 
+    Retrieval Augmented Generation (RAG) implementations with Google's <b>Gemma-2b-it</b> & <b>Gemini 1.0 Pro</b> 
+    models through HuggingFace and VertexAI respectively to generate travel recommendations. This early version (read 
+    quick and dirty implementation) aims to see if functionalities work smoothly. It relies on Wikipedia abstracts 
+    from 160 European cities to provide answers to your questions. Please be kind with it as it's a work in progress! 
+    </p> <br>Google Cloud credits are provided for this project. """)
+
+    with gr.Group():
+        query = gr.Textbox(label="Query", placeholder="Ask for your city recommendation here!")
+        model = gr.Dropdown(
+            ["google/gemma-2b-it", "gemini-1.0-pro"], label="Model", info="Select your model. Will add more models "
+                                                                          "later!",
+        )
+        output = gr.Textbox(label="Generated Results", lines=4)
+
+    with gr.Group():
+        with gr.Row():
+            submit_btn = gr.Button("Submit", variant="primary")
+            clear_btn = gr.Button("Clear", variant="secondary")
+            cancel_btn = gr.Button("Cancel", variant="stop")
+    submit_btn.click(generate_text, inputs=[query, model], outputs=[output])
+    clear_btn.click(clear, inputs=[], outputs=[query, model, output])
+    cancel_btn.click(clear, inputs=[], outputs=[query, model, output])
+
+    gr.Markdown("## Examples")
+    gr.Examples(
+        examples, inputs=[query, model], label="Examples", fn=generate_text, outputs=[output],
+        cache_examples=True,
+    )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(show_api=False)