Update app.py

app.py

@@ -9,6 +9,53 @@ from streaming import stream_to_gradio
 from huggingface_hub import login
 from gradio.data_classes import FileData
 
+import google.generativeai as genai
+os.environ["API_KEY"] = 'AIzaSyB8Hj4oCbBH9arFWSgybHnbpZLs2sa4p1w'
+os.environ["GOOGLE_API_KEY"] = 'AIzaSyBjuYTWBlHg4W2wGaQCKKbigz6deZuLUJc'
+genai.configure(api_key=os.environ["API_KEY"])
+
+generation_config = {
+  "temperature": 0.2,
+  "top_p": 0.95,
+  "top_k": 0,
+  "max_output_tokens": 8192,
+}
+
+safety_settings = [
+  {
+    "category": "HARM_CATEGORY_HARASSMENT",
+    "threshold": "BLOCK_MEDIUM_AND_ABOVE"
+  },
+  {
+    "category": "HARM_CATEGORY_HATE_SPEECH",
+    "threshold": "BLOCK_MEDIUM_AND_ABOVE"
+  },
+  {
+    "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+    "threshold": "BLOCK_MEDIUM_AND_ABOVE"
+  },
+  {
+    "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
+    "threshold": "BLOCK_MEDIUM_AND_ABOVE"
+  },
+]
+
+context = "You are an expert data analyst who can provide guidance around what needs to be analyzed from a dataset by just looking at metadata."
+system_instruction = context
+import re
+
+model = genai.GenerativeModel(model_name="gemini-1.5-pro-latest",
+                              generation_config=generation_config,
+                              system_instruction=system_instruction,
+                              safety_settings=safety_settings)
+
+
+def model_response(text):
+    #model = genai.GenerativeModel('gemini-pro')
+    response = model.generate_content(text)
+    return response.text
+
+
 login(os.getenv("HUGGINGFACEHUB_API_TOKEN"))
 
 llm_engine = HfEngine("meta-llama/Meta-Llama-3.1-70B-Instruct")
@@ -61,10 +108,14 @@ def interact_with_agent(file_input, additional_notes):
     - Columns with dtypes:
     {data_file.dtypes}"""
 
+    enhanced_notes = model_response(f'''Given the metadata of the dataset {data_structure_notes} and the context provided by the user {additional_notes}, figure out the
+    domain this dataset belongs to. Now assume the role of an expert data analyst in this domain and generate instructions/commentary that will help a large language model analyze
+    this dataset.''')
+
     prompt = base_prompt.format(structure_notes=data_structure_notes)
 
     if additional_notes and len(additional_notes) > 0:
-        prompt += "\nAdditional notes on the data:\n" + additional_notes
+        prompt += "\nAdditional notes on the data:\n" + enhanced_notes
 
     messages = [gr.ChatMessage(role="user", content=prompt)]
     yield messages + [