import gradio as gr
import cohere
import numpy as np
import warnings
from annoy import AnnoyIndex

# Cohere API Key 
cohere_api_key = "4GMQpN0nV9chIaV3DTQNmVYiskbTZdDBPrvw7pPs"
co = cohere.Client(cohere_api_key)

# Define a function to search for information
def search_text(query):
    query_embed = co.embed(texts=[query]).embeddings
    similar_item_ids = search_index.get_nns_by_vector(query_embed[0], 10, include_distances=True)
    search_results = texts[similar_item_ids[0]]
    return search_results

# Define a function to ask questions
def ask_llm(question, num_generations=1):
    results = search_text(question)
    context = results[0]
    prompt = f"""
    More information about Australian beaches at australia.com:
    {context}
    Question: {question}
    Extract the answer of the question from the text provided.
    If the text doesn't contain the answer,
    reply that the answer is not available."""
    prediction = co.generate(
        prompt=prompt,
        max_tokens=70,
        model="command-nightly",
        temperature=0.5,
        num_generations=num_generations
    )
    return prediction.generations

# Define your text data
question = "Which Sydney beach should I visit?"
text = """
Sydney is world famous for beautiful beaches. These beaches offer different vibes and attractions, from bustling crowds and great surfing conditions to more tranquil and family-friendly environments.
Bondi Beach: Bondi is perhaps the most famous beach in Sydney, if not Australia. It's known for its golden sands, vibrant atmosphere, and excellent surfing conditions. The Bondi to Coogee coastal walk offers stunning views of the coastline.
Manly Beach: Easily accessible by a scenic ferry ride from Circular Quay, Manly Beach is known for its laid-back atmosphere and family-friendly environment. The beach is great for swimming, surfing, and various water sports.
Cronulla Beach: Located in the Sutherland Shire, Cronulla offers a more relaxed atmosphere compared to some of the busier city beaches. It's a great spot for swimming, picnicking, and enjoying a range of seaside cafes and restaurants.
Bronte Beach: Situated between Bondi and Coogee, Bronte Beach is popular among both locals and visitors. It's a smaller, quieter beach with a beautiful park area and a natural rock pool that's ideal for swimming.
Tamarama Beach: Also known as "Glamarama" due to its popularity among the fashion-conscious crowd, Tamarama Beach is a smaller and more secluded option. It's surrounded by rocky cliffs and offers strong surf, making it a favorite among experienced surfers.
"""
texts = text.split('\n\n')
texts = np.array([t.strip(' \n') for t in texts if t])

# Create Gradio interface
iface = gr.Interface(
    fn=ask_llm,
    inputs=gr.inputs.Textbox(label="Ask a question about Sydney beaches:"),
    outputs="text",
    title="Sydney Beaches Information",
    live=True,
    capture_session=True,
)

# Launch the Gradio interface
iface.launch(share=True)