Create app.py

app.py

@@ -0,0 +1,122 @@
+from typing import Optional
+from sentence_transformers import SentenceTransformer
+import pymongo
+import sys
+from huggingface_hub import InferenceClient
+import gradio as gr
+
+sys.path.append("../")
+from config import constants
+
+HF_token = constants.HF_TOKEN
+
+
+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 = constants.MONGO_USERNAME
+    password = constants.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
+
+
+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],
+            ["Recommend places that are similar to Istanbul in terms of architecture", None],
+            ]
+
+demo = gr.Interface(
+    fn=generate_text,
+    inputs=["text",
+            gr.Dropdown(
+                ["google/gemma-2b-it","google/gemma-7b", "mistralai/Mixtral-8x7B-Instruct-v0.1"], label="Models", info="Will "
+                                                                                                             "add "
+                                                                                                             "more "
+                                                                                                             "models "
+                                                                                                             "later! "
+            ),
+            ],
+    title="🇪🇺 Euro TravelBot 🇪🇺",
+    description="Travel related queries for Europe.",
+    outputs=["text"],
+    examples=examples,
+)
+
+if __name__ == "__main__":
+    demo.launch()