test

app.py

@@ -1,6 +1,5 @@
 import gradio as gr
 import random
-import time
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # Load Vicuna 7B model and tokenizer
@@ -8,14 +7,6 @@ model_name = "lmsys/vicuna-7b-v1.3"
 model = AutoModelForCausalLM.from_pretrained(model_name)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-def respond_vicuna(message, chat_history, vicuna_chatbot):
-    input_ids = tokenizer.encode(message, return_tensors="pt")
-    output = model.generate(input_ids, max_length=50, num_beams=5, no_repeat_ngram_size=2)
-    bot_message = tokenizer.decode(output[0], skip_special_tokens=True)
-    chat_history.append((message, bot_message))
-    time.sleep(2)
-    return "", chat_history
-
 with gr.Blocks() as demo:
     gr.Markdown("# LLM Evaluator With Linguistic Scrutiny")
 
@@ -39,6 +30,7 @@ with gr.Blocks() as demo:
             prompt = gr.Textbox(show_label=False, placeholder="Enter prompt")
             send_button_POS = gr.Button("Send", scale=0)
         clear = gr.ClearButton([prompt, vicuna_chatbot1])
+    
     with gr.Tab("Chunk"):
         gr.Markdown("Strategy 1 QA")
         with gr.Row():
@@ -60,15 +52,12 @@ with gr.Blocks() as demo:
             send_button_Chunk = gr.Button("Send", scale=0)
         clear = gr.ClearButton([prompt_chunk, vicuna_chatbot1_chunk])
 
-    def respond(message, chat_history):
-        input_ids = tokenizer.encode(message, return_tensors="pt")
-        output = model.generate(input_ids, max_length=50, num_beams=5, no_repeat_ngram_size=2)
-        bot_message = tokenizer.decode(output[0], skip_special_tokens=True)
-        chat_history.append((message, bot_message))
-        time.sleep(2)
-        return "", chat_history
+    # Define the Gradio interface
+    def chatbot_interface(prompt):
+        vicuna_response = generate_response(model, tokenizer, prompt)
+        return {"Vicuna-7B": vicuna_response}
 
-    # Replace the old respond function with the new general function for Vicuna
-    prompt.submit(lambda message, chat_history: respond_vicuna(message, chat_history, vicuna_chatbot1), [prompt, vicuna_chatbot1, vicuna_chatbot1_chunk])
+    # Use the chatbot_interface function in the prompt.submit() method
+    prompt.submit(chatbot_interface, [prompt, vicuna_chatbot1, vicuna_chatbot1_chunk])
 
 demo.launch()