Update app.py

app.py

@@ -1,12 +1,14 @@
 import os
 from threading import Thread
 from typing import Iterator, List, Tuple
+import json
 
 import gradio as gr
 import spaces
 import torch
 import transformers
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from flask import Flask, request, jsonify
 
 DESCRIPTION = """\
 # Zero GPU Model Comparison Arena
@@ -28,6 +30,25 @@ MODEL_OPTIONS = [
 models = {}
 tokenizers = {}
 
+# Custom chat templates
+MISTRAL_TEMPLATE = """<s>[INST] {instruction} [/INST]
+{response}
+</s>
+<s>[INST] {instruction} [/INST]
+"""
+
+LLAMA_TEMPLATE = """<s>[INST] <<SYS>>
+You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.
+
+If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.
+<</SYS>>
+
+{instruction} [/INST]
+{response}
+</s>
+<s>[INST] {instruction} [/INST]
+"""
+
 for model_id in MODEL_OPTIONS:
     tokenizers[model_id] = AutoTokenizer.from_pretrained(model_id)
     models[model_id] = AutoModelForCausalLM.from_pretrained(
@@ -36,6 +57,22 @@ for model_id in MODEL_OPTIONS:
         load_in_8bit=True,
     )
     models[model_id].eval()
+    
+    # Set custom chat templates
+    if "Navarna" in model_id:
+        tokenizers[model_id].chat_template = MISTRAL_TEMPLATE
+    elif "OpenHathi" in model_id:
+        tokenizers[model_id].chat_template = LLAMA_TEMPLATE
+
+# Initialize Flask app
+app = Flask(__name__)
+
+@app.route('/log', methods=['POST'])
+def log_results():
+    data = request.json
+    # Here you can implement any additional processing or storage logic
+    print("Logged:", json.dumps(data, indent=2))
+    return jsonify({"status": "success"}), 200
 
 @spaces.GPU(duration=90)
 def generate(
@@ -101,11 +138,11 @@ def compare_models(
     chat_history1.append((message, output1))
     chat_history2.append((message, output2))
 
-    log_results(model1_name, model2_name, message, output1, output2)
+    log_comparison(model1_name, model2_name, message, output1, output2)
 
     return chat_history1, chat_history2, chat_history1, chat_history2
 
-def log_results(model1_name: str, model2_name: str, question: str, answer1: str, answer2: str, winner: str = None):
+def log_comparison(model1_name: str, model2_name: str, question: str, answer1: str, answer2: str, winner: str = None):
     log_data = {
         "question": question,
         "model1": {"name": model1_name, "answer": answer1},
@@ -113,12 +150,20 @@ def log_results(model1_name: str, model2_name: str, question: str, answer1: str,
         "winner": winner
     }
     
-    # Here you would implement the actual logging logic, e.g., sending to a server or writing to a file
-    print("Logged:", log_data)
+    # Send log data to Flask server
+    import requests
+    try:
+        response = requests.post('http://144.24.151.32:5000/log', json=log_data)
+        if response.status_code == 200:
+            print("Successfully logged to server")
+        else:
+            print(f"Failed to log to server. Status code: {response.status_code}")
+    except requests.RequestException as e:
+        print(f"Error sending log to server: {e}")
 
 def vote_better(model1_name, model2_name, question, answer1, answer2, choice):
     winner = model1_name if choice == "Model 1" else model2_name
-    log_results(model1_name, model2_name, question, answer1, answer2, winner)
+    log_comparison(model1_name, model2_name, question, answer1, answer2, winner)
     return f"You voted that {winner} performs better. This has been logged."
 
 with gr.Blocks(css="style.css") as demo:
@@ -166,4 +211,9 @@ with gr.Blocks(css="style.css") as demo:
     )
 
 if __name__ == "__main__":
-    demo.queue(max_size=10).launch()
+    # Start Flask server in a separate thread
+    flask_thread = Thread(target=app.run, kwargs={"host": "0.0.0.0", "port": 5000})
+    flask_thread.start()
+    
+    # Start Gradio app
+    demo.queue(max_size=10).launch()