Update app.py

app.py

@@ -26,10 +26,8 @@ sd_pipe = StableDiffusionPipeline.from_pretrained(
 def get_current_weather(place: str) -> str:
     """
     Get the current weather for a given location.
-
     Args:
         place (str): Name of the city or location (e.g., "London" or "New York").
-
     Returns:
         str: Weather condition, temperature, humidity, and wind speed.
     """
@@ -60,11 +58,9 @@ def get_current_weather(place: str) -> str:
 def get_current_time_in_timezone(timezone: str) -> str:
     """
     Get the current local time in a given timezone.
-
     Args:
         timezone (str): Timezone string in the format 'Region/City', 
                         e.g., "America/New_York".
-
     Returns:
         str: Formatted local time string.
     """
@@ -83,11 +79,9 @@ qa_pipeline = pipeline("text2text-generation", model="google/flan-t5-base")
 def document_qna_tool(pdf_path: str, question: str) -> str:
     """
     Answer a natural language question based on the content of a PDF document.
-
     Args:
         pdf_path (str): Path to the local PDF file.
         question (str): Natural language question to ask about the PDF.
-
     Returns:
         str: Answer generated from the most relevant section of the document.
     """
@@ -128,7 +122,6 @@ def image_generator(prompt: str) -> str:
 
 # -------------------- Local LLM (Replaces HfApiModel) --------------------
 from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
 
 class LocalModel:
     def __init__(self):
@@ -140,10 +133,25 @@ class LocalModel:
             device_map="auto" if torch.cuda.is_available() else None,
         )
 
-    def generate(self, prompt, max_new_tokens=500):
+    def generate(self, prompt, max_new_tokens=500, **kwargs):
+        """
+        Generate text from the given prompt.
+        Extra kwargs like 'stop_sequences' are accepted for compatibility.
+        """
+        stop_sequences = kwargs.pop("stop_sequences", None)
+
         inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
         output = self.model.generate(**inputs, max_new_tokens=max_new_tokens)
-        return self.tokenizer.decode(output[0], skip_special_tokens=True)
+        text = self.tokenizer.decode(output[0], skip_special_tokens=True)
+
+        # If stop_sequences provided, truncate output
+        if stop_sequences:
+            for stop in stop_sequences:
+                if stop in text:
+                    text = text.split(stop)[0]
+                    break
+
+        return text
 
     def __call__(self, prompt, **kwargs):
         return self.generate(prompt, **kwargs)