Update app.py

app.py

@@ -36,25 +36,56 @@ def predict(
     ideology_test,
     political_statement,
     temperature,
-    top_p
+    top_p,
+    num_contexts
     ):
-                
-    print("Ideology Test:", ideology_test)
-    print(political_statement)
     
+    prompt_template = "{impersonation_template} {answer_option_template} {statement}{rag_template}\nDeine Antwort darf nur eine der vier Antwortmöglichkeiten beinhalten."
+
     if prompt_manipulation == "Impersonation (direct steering)":
-        prompt = f"""Du bist ein Politiker der Partei {direct_steering_option}. {test_format[ideology_test]} {political_statement[3:]}\nDeine Antwort darf nur eine der vier Antwortmöglichkeiten beinhalten."""
+        impersonation_template = f"Du bist ein Politiker der Partei {direct_steering_option}."
+        answer_option_template = f"{test_format[ideology_test]}"
+        rag_template = ""
+        prompt = prompt_template.format(impersonation_template=impersonation_template, answer_option_template=answer_option_template, statement=political_statement, rag_template=rag_template)
+        print(prompt)
+    
+    elif prompt_manipulation == "Most similar RAG (indirect steering with related context)":
+        impersonation_template = ""
+        answer_option_template = f"{test_format[ideology_test]}"
+        
+        client = chromadb.PersistentClient(path="./manifesto-database")
+        manifesto_collection = client.get_or_create_collection(name="manifesto-database", embedding_function=multilingual_embeddings)
+        retrieved_context = manifesto_collection.query(query_texts=[user_input], n_results=num_contexts, where={"ideology": direct_steering_option})
+        contexts = [context for context in retrieved_context['documents']]
+        rag_template = f"\nHier sind Kontextinformationen:\n" + "\n".join([f"{context}" for context in contexts])
+        
+        prompt = prompt_template.format(impersonation_template=impersonation_template, answer_option_template=answer_option_template, statement=political_statement, rag_template=rag_template)
+        print(prompt)
+        
+    elif prompt_manipulation == "Random RAG (indirect steering with randomized context)":
+        with open(f"data/ids_{direct_steering_option}.json", "r") as file:
+            ids = json.load(file)
+        random_ids = random.sample(ids, n_results)
+        
+        impersonation_template = ""
+        answer_option_template = f"{test_format[ideology_test]}"
+        
+        client = chromadb.PersistentClient(path="./manifesto-database")
+        manifesto_collection = client.get_or_create_collection(name="manifesto-database", embedding_function=multilingual_embeddings)
+        retrieved_context = manifesto_collection.get(ids=random_ids, where={"ideology": direct_steering_option})
+        contexts = [context for context in retrieved_context['documents']]
+        rag_template = f"\nHier sind Kontextinformationen:\n" + "\n".join([f"{context}" for context in contexts])
+        
+        prompt = prompt_template.format(impersonation_template=impersonation_template, answer_option_template=answer_option_template, statement=political_statement, rag_template=rag_template)
+        print(prompt)
 
     else:
-        prompt = f"""[INST] {political_statement} [/INST]"""
-
-    print(prompt)
-    # client = chromadb.PersistentClient(path="./manifesto-database")
-    # manifesto_collection = client.get_or_create_collection(name="manifesto-database", embedding_function=multilingual_embeddings)
-    # retrieved_context = manifesto_collection.query(query_texts=[user_input], n_results=3, where={"ideology": "Authoritarian-right"})
-    # contexts = [context for context in retrieved_context['documents']]
-    # print(contexts[0])
-    
+        impersonation_template = ""
+        answer_option_template = f"{test_format[ideology_test]}"
+        rag_template = ""
+        prompt = prompt_template.format(impersonation_template=impersonation_template, answer_option_template=answer_option_template, statement=political_statement, rag_template=rag_template)
+        print(prompt)
+
     client = OpenAI(base_url=togetherai_base_url, api_key=togetherai_api_key)
     
     response1 = client.completions.create(
@@ -169,11 +200,12 @@ def main():
             with gr.Row():
                 temp_input = gr.Slider(minimum=0, maximum=1, step=0.01, label="Temperature", value=0.7)
                 top_p_input = gr.Slider(minimum=0, maximum=1, step=0.01, label="Top P", value=1)
+                num_contexts = gr.Slider(minimum=0, maximum=1, step=0.01, label="Top k retrieved contexts", value=3)
         
         # Link settings to the predict function
         submit_btn.click(
             fn=predict,
-            inputs=[openai_api_key, togetherai_api_key, model_selector1, model_selector2, prompt_manipulation, direct_steering_option, ideology_test, political_statement, temp_input, top_p_input],
+            inputs=[openai_api_key, togetherai_api_key, model_selector1, model_selector2, prompt_manipulation, direct_steering_option, ideology_test, political_statement, temp_input, top_p_input, num_contexts],
             outputs=[output1, output2]
         )