Update app.py

app.py

@@ -1,37 +1,78 @@
-import gradio as gr
 import pip
 import os
 pip.main(['install', 'transformers'])
 pip.main(['install', 'torch'])
 pip.main(['install', 'pymongo'])
-import pymongo
+import gradio as gr
 from transformers import pipeline
-from datetime import datetime
-model_name_or_path='dgnk007/eagle'
-generate=pipeline('text-generation',model=model_name_or_path)
-
-myclient = pymongo.MongoClient(os.environ['DB_URI'])
-mydb = myclient["eagle"]
-def store_at_db(prompt,response):
-    rawusage = mydb["rawusage"]
-    rawusage.insert_one({"prompt":prompt,"response":response,"current_time":datetime.now()})
-
-def generate_response(message):
-    prompt_template=f"Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\n {message}\n\n### Response:\n"
-    response=generate(prompt_template,max_length=1024,return_full_text=False,eos_token_id=21017,pad_token_id=50256)
-    store_at_db(message,response[0]['generated_text'])
-    return response[0]['generated_text']
-
-examples = [
-    ["Give three tips for staying healthy."],
-    ["How can we reduce air pollution?"],
-    ["Explain what an API is."],
-]
-demo = gr.Interface(
-    fn=generate_response,
-    inputs="text",
-    outputs="text",
-    examples=examples,
-    #flagging_callback=gr.HuggingFaceDatasetSaver(hf_token=os.environ['DatasetKey'],dataset_name=os.environ['DatasetName'],private=True,verbose=True )
-)
-demo.launch()
+import pymongo
+
+mongo_client = pymongo.MongoClient(os.environ['DB_URI'])
+db = mongo_client["gradio_db"]
+btn_disable=gr.Button.update(interactive=False)
+btn_enable=gr.Button.update(interactive=True)
+generator = pipeline("text-generation", model="dgnk007/eagle2")
+
+def store_in_mongodb(collection_name, data):
+    collection = db[collection_name]
+    return collection.insert_one(data)
+    
+
+def generate_text(message,sequences):
+    prompt_template=f"Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\n {message}\n\n### Response:\n\n"
+    generated_text = generator(prompt_template, max_length=1024,return_full_text=False,eos_token_id=21017,pad_token_id=50256, num_return_sequences=sequences)
+    return generated_text
+
+def general_function(input_text):
+    output_text = generate_text(input_text,1)[0]['generated_text']
+    store_in_mongodb("general_collection", {"input": input_text, "output": output_text})
+    return  output_text
+
+def arena_function(input_text):
+    output_text = generate_text(input_text,2)
+    data_to_store = {
+        "input": input_text,
+        "r1": output_text[0]['generated_text'],
+        "r2": output_text[1]['generated_text'],
+    }
+    id=store_in_mongodb("arena_collection", data_to_store)
+    return  output_text[0]['generated_text'], output_text[1]['generated_text'], id.inserted_id,btn_enable,btn_enable,btn_enable,btn_enable
+
+general_interface = gr.Interface(fn=general_function, inputs=gr.Textbox(label="Enter your text here:", min_width=600), outputs="text")
+
+def reward_click(id,reward):
+    db["arena_collection"].update_one(
+        {"_id": id},
+        {"$set": {"reward": reward}}
+    )
+    return btn_disable,btn_disable,btn_disable,btn_disable
+
+with gr.Blocks() as arena_interface:
+    obid=gr.State([])
+    with gr.Row():
+        with gr.Column():
+            input_box = gr.Textbox(label="Enter your text here:", min_width=600)
+            prompt = gr.Button("Submit", variant="primary")
+    with gr.Row():
+        gr.Examples(['what is google?','what is youtube?'], input_box,)
+    with gr.Row():
+        output_block = [
+            gr.Textbox(label="Response 1", interactive=False),
+            gr.Textbox(label="Response 2", interactive=False),
+            obid
+        ]
+    with gr.Row():
+        tie=gr.Button(value="Tie",size='sm',interactive=False)
+        r1=gr.Button(value="Response 1 Wins",variant='primary',interactive=False)
+        r2=gr.Button(value="Response 2 Wins",variant='primary',interactive=False)
+        bad=gr.Button(value="Both are Bad",variant='secondary',interactive=False)
+        buttonGroup=[tie,r1,r2,bad]
+    prompt.click(fn=arena_function, inputs=input_box, outputs=output_block+buttonGroup)
+    tie.click(fn=reward_click,inputs=[obid,gr.State('tie')],outputs=buttonGroup)
+    r1.click(fn=reward_click,inputs=[obid,gr.State('r1')],outputs=buttonGroup)
+    r2.click(fn=reward_click,inputs=[obid,gr.State('r2')],outputs=buttonGroup)
+    bad.click(fn=reward_click,inputs=[obid,gr.State('bad')],outputs=buttonGroup)
+demo = gr.TabbedInterface([general_interface, arena_interface], ["General", "Arena"])
+
+
+demo.launch()