Update modules.py

modules.py

@@ -1,108 +1,110 @@
-import time
-from typing import Any, List, Optional, Tuple, Union
-import numpy as np
-import pandas as pd
-import pickle
-import glob
-from pydantic import SecretStr
-from langchain_community.llms.ollama import Ollama
-from langchain_community.embeddings import OllamaEmbeddings
-from langchain_google_genai.llms import GoogleGenerativeAI
-from langchain_google_genai.embeddings import GoogleGenerativeAIEmbeddings
-import streamlit as st
-
-
-# Streamlit functions to select the model
-def model_selection()->Union[str, Tuple[str, SecretStr]]:
-    select = st.sidebar.selectbox("Select model", options=["Ollama", "Gemini"])
-    if select == "Ollama":
-        return select, None
-    elif select == "Gemini":
-        api_key = st.sidebar.text_input("Enter API key", type="password")
-        return select, api_key
-    else:
-        raise ValueError("Invalid model name. Please choose 'Gemini' or 'Ollama'.")
-    
-
-# This function will be used to time the functions
-def timer(func)->callable:
-    def wrapper(*args, **kwargs):
-        start = time.perf_counter()
-        result = func(*args, **kwargs)
-        end = time.perf_counter()
-        print(f"{func.__name__} took {end - start} seconds")
-        return result
-    return wrapper
-
-@timer
-def text2embeddings(document: Union[str, List[str]], api_key:str, model:str) -> np.ndarray:
-    google_embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001", google_api_key=api_key)
-    ollama_embeddings = OllamaEmbeddings(model="phi3")
-    
-    if model == "Gemini":
-        if isinstance(document, str):
-            return np.array(google_embeddings.embed_query(document))
-        return np.array(google_embeddings.embed_documents(document))
-    elif model == "Ollama":
-        if isinstance(document, str):
-            return np.array(ollama_embeddings.embed_query(document))
-        return np.array(ollama_embeddings.embed_documents(document))
-    else:
-        raise ValueError("Invalid model name. Please choose 'Gemini' or 'Ollama'.")
-
-
-@timer
-def build_context_matrix(df:str, column:str, api_key:str, model:str)->Any:
-    dataframe = pd.read_csv(df)
-    if glob.glob("Data/context_matrix.pkl"):
-        pickle_file = pickle.load(open("Data/context_matrix.pkl", "rb"))
-        return pickle_file 
-    else:
-        matrix = text2embeddings(document=dataframe[column].tolist(), api_key=api_key, model=model)
-        pickle.dump(matrix, open("Data/context_matrix.pkl", "wb"))
-        print("Context matrix created, Run the code again")
-        return None
-    
-
-@timer
-def semantic_chunk(query:str, df:str, column:str, api_key:str, model:str, chunk_size:int)->str:
-    dataframe = pd.read_csv(df)
-    matrix = build_context_matrix(df=df, column=column, api_key=api_key, model=model)
-    query_vector = text2embeddings(document=query, api_key=api_key, model=model)
-    score = matrix @ query_vector
-    top_k = np.argsort(score)[::-1][:chunk_size]
-    context = dataframe[column].iloc[top_k]
-    string = "\n".join(context)
-    return string
-
-
-def select_llm(model:str, api_key:str=None, temperature: float=0.5)->Union[Ollama, GoogleGenerativeAI]:
-    if model == "Ollama":
-        return Ollama(model="phi3", temperature=temperature)
-    elif model == "Gemini":
-        return GoogleGenerativeAI(model="gemini-pro", temperature=temperature, google_api_key=api_key)
-    else:
-        raise ValueError("Invalid model name. Please choose 'Gemini' or 'Ollama'.")
-
-
-def get_prompt(context:str, query:str)->str:
-    prompt = f"""You are a chatbot named Mira. You answer user's query about ecommerce products.
-    Here are some rules you will follow:
-    1. Your response will be concise and informative.
-    2. You must answer the user's query from the given context.
-    3. Do not generate any incomplete sentences or duplicate sentences.
-    
-    -----------------------------------------------
-    Context: '''{context}'''
-    -----------------------------------------------
-    Question: '''{query}'''
-    -----------------------------------------------
-    """
-    return prompt
-
-def retrieve_query(query:str, df:str, column:str, api_key:str, model:str, temperature:float=0.5, chunk_size:int=100)->str:
-    context = semantic_chunk(query=query, df=df, column=column, model=model, api_key=api_key, chunk_size=chunk_size)
-    llm = select_llm(model=model, api_key=api_key, temperature=temperature)
-    prompt = get_prompt(context, query)
-    response = llm.stream(prompt)
+import time
+from typing import Any, List, Optional, Tuple, Union
+import numpy as np
+import pandas as pd
+import pickle
+import glob
+from pydantic import SecretStr
+from langchain_community.llms.ollama import Ollama
+from langchain_community.embeddings import OllamaEmbeddings
+from langchain_google_genai.llms import GoogleGenerativeAI
+from langchain_google_genai.embeddings import GoogleGenerativeAIEmbeddings
+import streamlit as st
+
+
+# Streamlit functions to select the model
+def model_selection()->Union[str, Tuple[str, SecretStr]]:
+    select = st.sidebar.selectbox("Select model", options=["Ollama", "Gemini"])
+    if select == "Ollama":
+        return select, None
+    elif select == "Gemini":
+        api_key = st.sidebar.text_input("Enter API key", type="password")
+        return select, api_key
+    else:
+        raise ValueError("Invalid model name. Please choose 'Gemini' or 'Ollama'.")
+    
+
+# This function will be used to time the functions
+def timer(func)->callable:
+    def wrapper(*args, **kwargs):
+        start = time.perf_counter()
+        result = func(*args, **kwargs)
+        end = time.perf_counter()
+        print(f"{func.__name__} took {end - start} seconds")
+        return result
+    return wrapper
+
+@timer
+def text2embeddings(document: Union[str, List[str]], api_key:str, model:str) -> np.ndarray:
+    google_embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001", google_api_key=api_key)
+    ollama_embeddings = OllamaEmbeddings(model="phi3")
+    
+    if model == "Gemini":
+        if isinstance(document, str):
+            return np.array(google_embeddings.embed_query(document))
+        return np.array(google_embeddings.embed_documents(document))
+    elif model == "Ollama":
+        if isinstance(document, str):
+            return np.array(ollama_embeddings.embed_query(document))
+        return np.array(ollama_embeddings.embed_documents(document))
+    else:
+        raise ValueError("Invalid model name. Please choose 'Gemini' or 'Ollama'.")
+
+
+@timer
+def build_context_matrix(df:str, column:str, api_key:str, model:str)->Any:
+    dataframe = pd.read_csv(df)
+    if glob.glob("Data/context_matrix.pkl"):
+        pickle_file = pickle.load(open("Data/context_matrix.pkl", "rb"))
+        return pickle_file 
+    else:
+        matrix = text2embeddings(document=dataframe[column].tolist(), api_key=api_key, model=model)
+        pickle.dump(matrix, open("Data/context_matrix.pkl", "wb"))
+        print("Context matrix created, Run the code again")
+        return None
+    
+
+@timer
+def semantic_chunk(query:str, df:str, column:str, api_key:str, model:str, chunk_size:int)->str:
+    dataframe = pd.read_csv(df)
+    matrix = build_context_matrix(df=df, column=column, api_key=api_key, model=model)
+    query_vector = text2embeddings(document=query, api_key=api_key, model=model)
+    score = matrix @ query_vector
+    top_k = np.argsort(score)[::-1][:chunk_size]
+    context = dataframe[column].iloc[top_k]
+    string = "\n".join(context)
+    return string
+
+
+def select_llm(model:str, api_key:str=None, temperature: float=0.5)->Union[Ollama, GoogleGenerativeAI]:
+    if model == "Ollama":
+        return Ollama(model="phi3", temperature=temperature)
+    elif model == "Gemini":
+        return GoogleGenerativeAI(model="gemini-pro", temperature=temperature, google_api_key=api_key)
+    else:
+        raise ValueError("Invalid model name. Please choose 'Gemini' or 'Ollama'.")
+
+
+def get_prompt(context:str, query:str)->str:
+    prompt = f"""You are a chatbot named Mira. You answer user's query about ecommerce products.
+    Here are some rules you will follow:
+    1. Your response will be concise and informative.
+    2. You must answer the user's query from the given context.
+    3. Do not generate any incomplete sentences or duplicate sentences.
+    
+    -----------------------------------------------
+    Context: '''{context}'''
+    -----------------------------------------------
+    Question: '''{query}'''
+    -----------------------------------------------
+    """
+    return prompt
+
+def retrieve_query(query:str, df:str, column:str, api_key:str, model:str, temperature:float=0.5, chunk_size:int=100)->str:
+    if not api_key:
+        st.write("Please enter your key")
+    context = semantic_chunk(query=query, df=df, column=column, model=model, api_key=api_key, chunk_size=chunk_size)
+    llm = select_llm(model=model, api_key=api_key, temperature=temperature)
+    prompt = get_prompt(context, query)
+    response = llm.stream(prompt)
     return response