Update app.py

app.py

@@ -1,31 +1,36 @@
-# Import required libraries
-import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import gradio as gr
+import torch
 
-# Load the model and tokenizer
+# Define model name
 MODEL_NAME = "SeaLLMs/SeaLLM-7B-v2.5"
 
-# Download model and tokenizer from Hugging Face
+# Load the model and tokenizer with optimized settings
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype="auto", device_map="auto")
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    torch_dtype=torch.float16,  # Use float16 for GPU optimization
+    device_map="auto"          # Automatically assign to available GPUs
+)
 
-# Define the chatbot function
-def chatbot(user_input):
-    inputs = tokenizer(user_input, return_tensors="pt").to("cuda")
-    outputs = model.generate(inputs["input_ids"], max_length=150, num_return_sequences=1, temperature=0.7)
+# Chatbot function
+def chatbot(prompt):
+    # Tokenize input
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+    # Generate response
+    outputs = model.generate(inputs.input_ids, max_new_tokens=150, temperature=0.7)
+    # Decode and return response
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return response
 
-# Create a Gradio interface
-interface = gr.Interface(
+# Gradio Interface
+iface = gr.Interface(
     fn=chatbot,
-    inputs="text",
-    outputs="text",
+    inputs=gr.Textbox(label="Ask me anything:", lines=3, placeholder="Type your message here..."),
+    outputs=gr.Textbox(label="Response"),
     title="SeaLLM Chatbot",
-    description="A chatbot powered by SeaLLM-7B-v2.5.",
-    examples=["Hello!", "What's the weather today?", "Tell me a joke!"],
+    description="A chatbot powered by SeaLLM-7B-v2.5 for text generation.",
 )
 
-# Launch the interface
 if __name__ == "__main__":
-    interface.launch()
+    iface.launch()