Update app.py

app.py

@@ -2,11 +2,11 @@ import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Load the Qwen2.5-72B-Instruct model
-model_name = "Qwen/Qwen2.5-72B-Instruct"
+# Load the Zephyr-7B model
+model_name = "HuggingFaceH4/zephyr-7b-beta"
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
-    torch_dtype="auto",
+    torch_dtype=torch.bfloat16,
     device_map="auto"
 )
 tokenizer = AutoTokenizer.from_pretrained(model_name)
@@ -37,34 +37,29 @@ def chat(user_input):
         # Append user message to the conversation history
         messages.append({"role": "user", "content": user_input})
 
-        # Prepare input for the model
-        text = tokenizer.apply_chat_template(
-            messages,
-            tokenize=False,
-            add_generation_prompt=True
-        )
-        model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+        # Prepare input for the model using chat template
+        chat_input = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
 
         # Generate a response using the model
         try:
+            model_inputs = tokenizer(chat_input, return_tensors="pt").to(model.device)
             generated_ids = model.generate(
                 **model_inputs,
-                max_new_tokens=512
+                max_new_tokens=256,
+                do_sample=True,
+                temperature=0.7,
+                top_k=50,
+                top_p=0.95
             )
-            generated_ids = [
-                output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
-            ]
-
-            # Decode the response
-            response_content = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
+            response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
 
         except Exception as e:
-            response_content = f"Error: {str(e)}"
+            response = f"Error: {str(e)}"
 
         # Store assistant response in the chat history
-        messages.append({"role": "assistant", "content": response_content})
+        messages.append({"role": "assistant", "content": response})
 
-        return messages, response_content
+        return messages, response
     return messages, ""
 
 # Gradio Interface
@@ -96,12 +91,12 @@ with gr.Blocks() as demo:
     user_input = gr.Textbox(label="Your Message", placeholder="Type your message here...")
 
     # Connect the buttons to their respective functions
-    output_message = gr.Textbox(label="Output Message", interactive=False)
+    output_message = gr.Textbox(label="Output Message")
     submit_btn.click(submit_questionnaire, inputs=[name, age, location, gender, ethnicity, height, weight,
                                                     style_preference, color_palette, everyday_style], outputs=output_message)
 
-    reset_btn.click(reset_chat, outputs=[chatbox, output_message])  # Corrected outputs
-    user_input.submit(chat, inputs=user_input, outputs=[chatbox, user_input])  # Fixed to include chatbox output
+    reset_btn.click(reset_chat, outputs=[chatbox])  # Reset chat
+    user_input.submit(chat, inputs=user_input, outputs=[chatbox, user_input])  # Connect chat input
 
 # Run the app
 demo.launch()