test

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import random
+import time
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # Load Vicuna 7B model and tokenizer
@@ -30,7 +31,6 @@ 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():
@@ -52,12 +52,21 @@ with gr.Blocks() as demo:
             send_button_Chunk = gr.Button("Send", scale=0)
         clear = gr.ClearButton([prompt_chunk, vicuna_chatbot1_chunk])
 
+    # Define the function for generating responses
+    def generate_response(model, tokenizer, prompt):
+        inputs = tokenizer(prompt, return_tensors="pt")
+        outputs = model.generate(**inputs, max_length=500, pad_token_id=tokenizer.eos_token_id)
+        response = tokenizer.decode(outputs[0])
+        return response
+
     # Define the Gradio interface
     def chatbot_interface(prompt):
         vicuna_response = generate_response(model, tokenizer, prompt)
+        # llama_response = generate_response(llama_model, llama_tokenizer, prompt)
+
         return {"Vicuna-7B": vicuna_response}
 
-    # Use the chatbot_interface function in the prompt.submit() method
+    # Replace the old respond function with the new general function for Vicuna
     prompt.submit(chatbot_interface, [prompt, vicuna_chatbot1, vicuna_chatbot1_chunk])
 
 demo.launch()