Update app.py

app.py

@@ -72,8 +72,19 @@ def generate_response(text: str, model_gemma, tokenizer_gemma, device) -> str: #
         )
 
         generated_output = model_gemma.generate(input, generation_config=generation_config)
-        return tokenizer_gemma.decode(generated_output[0], skip_special_tokens=True)
-        
+        decoded_output = tokenizer_gemma.decode(generated_output[0], skip_special_tokens=False)
+
+        # Extract the assistant's response (Gemma specific)
+        start_token = "<start_of_turn>model"
+        end_token = "<end_of_turn>"
+
+        start_index = decoded_output.find(start_token)
+        if start_index != -1:
+            start_index += len(start_token)
+            end_index = decoded_output.find(end_token, start_index)
+            assistant_response = decoded_output[start_index:].strip()
+            return assistant_response
+        return decoded_output
         #input_text = "Reapond to the users prompt: " + text
         #input = tokenizer_gemma(input_text, return_tensors="pt").to(device)
         #generated_output = model_gemma.generate(**input, max_length=MAX_GEMMA_LENGTH, early_stopping=True)