Update app.py

app.py

@@ -24,6 +24,9 @@ class CustomTransformersModel:
         self.model = AutoModelForCausalLM.from_pretrained(model_id)
         
     def __call__(self, prompt, **kwargs):
+        # Extract and handle stop_sequences if present
+        stop_sequences = kwargs.pop('stop_sequences', None)
+        
         # Format the prompt using our chat template
         messages = [{"role": "user", "content": prompt}]
         formatted_prompt = self.tokenizer.apply_chat_template(messages, tokenize=False)
@@ -43,6 +46,12 @@ class CustomTransformersModel:
         # Decode the response
         response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
         
+        # Apply stop sequences manually if provided
+        if stop_sequences:
+            for stop_seq in stop_sequences:
+                if stop_seq in response:
+                    response = response.split(stop_seq)[0]
+        
         # Extract just the assistant's response
         try:
             assistant_response = response.split("Assistant: ")[-1]