Update app.py

app.py

@@ -2,10 +2,22 @@ from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from huggingface_hub import login
+from transformers import StoppingCriteria, StoppingCriteriaList
 import os
 import torch
 import uvicorn
 
+class StopOnStrings(StoppingCriteria):
+    def __init__(self, tokenizer, stop_strings):
+        self.tokenizer = tokenizer
+        self.stop_ids = [tokenizer.encode(s, add_special_tokens=False) for s in stop_strings]
+
+    def __call__(self, input_ids, scores, **kwargs):
+        for stop_id in self.stop_ids:
+            if input_ids[0][-len(stop_id):].tolist() == stop_id:
+                return True
+        return False
+
 login(os.getenv("HF_TOKEN"))
 
 app = FastAPI(
@@ -36,13 +48,18 @@ async def generate_text(request: GenerateRequest):
         inputs = tokenizer(request.prompt, return_tensors="pt").to(model.device)
         
         with torch.no_grad():
+            stopping = StoppingCriteriaList([
+                StopOnStrings(tokenizer, ["\n\n", "###", "END"])
+            ])
+            
             outputs = model.generate(
                 **inputs,
                 max_new_tokens=request.max_new_tokens,
                 temperature=request.temperature,
                 do_sample=True,
                 repetition_penalty=1.1,
-                pad_token_id=tokenizer.eos_token_id
+                pad_token_id=tokenizer.eos_token_id,
+                stopping_criteria=stopping
             )
         
         full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)