Update app.py

app.py

@@ -24,26 +24,35 @@ def load_model():
     return tokenizer, model
 
 tokenizer, model = load_model()
+# Function to generate chatbot response using the provided template
+def get_completion(query: str, model, tokenizer) -> str:
+    device = "cuda:0" if torch.cuda.is_available() else "cpu" #Use cuda if available.
 
-# Function to generate chatbot response
-def generate_response(prompt, chat_history=""):
-    inputs = tokenizer.encode(chat_history + prompt, return_tensors="pt")
+    prompt_template = f"""
+<start_of_turn>system You are a support chatbot who helps with user queries chatbot who always responds in the style of a professional.\n<end_of_turn>
+<start_of_turn>user
 
-    # Generate a response
-    outputs = model.generate(
-        inputs,
-        max_length=1000,
-        pad_token_id=tokenizer.eos_token_id,
-        temperature=0.7,
-        top_k=50,
-        top_p=0.95,
-    )
 
-    response = tokenizer.decode(outputs[0][inputs.shape[-1]:], skip_special_tokens=True)
-    return response
+{query}
+<end_of_turn>
+
+<start_of_turn>model
+"""
+    prompt = prompt_template.format(query=query)
+
+    encodeds = tokenizer(prompt, return_tensors="pt", add_special_tokens=True)
+
+    model_inputs = encodeds.to(device)
+
+    model.to(device) #Move model to device.
+
+    generated_ids = model.generate(**model_inputs, max_new_tokens=1000, do_sample=True, pad_token_id=tokenizer.eos_token_id)
+    decoded = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+    model_response = decoded.split("model\n")[-1].strip()
+    return model_response
 
 # Streamlit app
-st.title("Gemma-2b-it Chatbot")
+st.title("Customer Care Chatbot")
 
 # Initialize chat history
 if "messages" not in st.session_state: