Update app.py

app.py

@@ -6,45 +6,28 @@ import torch
 # Load base model
 base_model = AutoModelForCausalLM.from_pretrained(
     "mistralai/Mistral-7B-Instruct-v0.1",
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    torch_dtype=torch.float16,
     device_map="auto"
 )
 
-# Load LoRA adapter
+# Load LoRA Adapter
 model = PeftModel.from_pretrained(base_model, "gaurav2003/room-service-chatbot")
 
-# Load tokenizer
-tokenizer = AutoTokenizer.from_pretrained("gaurav2003/room-service-chatbot")
-tokenizer.pad_token = tokenizer.eos_token  # Set pad token if missing
+# Load tokenizer (from base model)
+tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1")
+tokenizer.pad_token = tokenizer.eos_token
 
-# Chat function
 def chat(user_input, history=[]):
-    input_ids = tokenizer(user_input, return_tensors="pt").input_ids.to(model.device)
-
-    # Generate response
-    with torch.no_grad():
-        output_ids = model.generate(
-            input_ids,
-            max_new_tokens=100,
-            pad_token_id=tokenizer.pad_token_id,
-            do_sample=True,
-            top_p=0.95,
-            temperature=0.7
-        )
-
-    # Decode response (remove the prompt part)
-    response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    response = response.replace(user_input, "").strip()
-
+    input_ids = tokenizer(user_input, return_tensors="pt", padding=True).input_ids.to(model.device)
+    output = model.generate(input_ids, max_new_tokens=150)
+    response = tokenizer.decode(output[0], skip_special_tokens=True)
     return response
 
-# Gradio UI
 iface = gr.Interface(
     fn=chat,
-    inputs=gr.Textbox(placeholder="Ask something...", label="User Input"),
+    inputs=gr.Textbox(placeholder="Ask something..."),
     outputs="text",
-    title="Room Service Chatbot",
-    description="Chat with your fine-tuned hotel assistant!"
+    title="Room Service Chatbot"
 )
 
 if __name__ == "__main__":