Update app.py

app.py

@@ -9,7 +9,8 @@ from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
     GenerationConfig,
-    StoppingCriteriaList
+    StoppingCriteriaList,
+    StoppingCriteria
 )
 import uvicorn
 import asyncio
@@ -48,9 +49,11 @@ class GenerateRequest(BaseModel):
 
 class LocalModelLoader:
     def __init__(self):
-        pass
+        self.loaded_models = {}
 
     async def load_model_and_tokenizer(self, model_name):
+        if model_name in self.loaded_models:
+            return self.loaded_models[model_name]
         try:
             config = AutoConfig.from_pretrained(model_name)
             tokenizer = AutoTokenizer.from_pretrained(model_name, config=config)
@@ -59,12 +62,24 @@ class LocalModelLoader:
             if tokenizer.eos_token_id is not None and tokenizer.pad_token_id is None:
                 tokenizer.pad_token_id = config.pad_token_id or tokenizer.eos_token_id
 
+            self.loaded_models[model_name] = (model, tokenizer)
             return model, tokenizer
         except Exception as e:
             raise HTTPException(status_code=500, detail=f"Error loading model: {e}")
 
 model_loader = LocalModelLoader()
 
+class StopOnTokens(StoppingCriteria):
+    def __init__(self, stop_token_ids: list[int]):
+        self.stop_token_ids = stop_token_ids
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        for stop_id in self.stop_token_ids:
+            if input_ids[0][-1] == stop_id:
+                return True
+        return False
+
+
 @app.post("/generate")
 async def generate(request: GenerateRequest):
     try:
@@ -96,69 +111,40 @@ async def generate(request: GenerateRequest):
             num_return_sequences=num_return_sequences,
         )
 
+        stop_token_ids = []
+        if stop_sequences:
+            stop_token_ids = tokenizer.convert_tokens_to_ids(stop_sequences)
+        stopping_criteria_list = StoppingCriteriaList([StopOnTokens(stop_token_ids)]) if stop_token_ids else None
+
+
         return StreamingResponse(
-            stream_text(model, tokenizer, input_text, generation_config, stop_sequences, device, chunk_delay),
+            stream_text(model, tokenizer, input_text, generation_config, stopping_criteria_list, device, chunk_delay),
             media_type="text/plain"
         )
 
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
 
-async def stream_text(model, tokenizer, input_text, generation_config, stop_sequences, device, chunk_delay, max_length=2048):
+async def stream_text(model, tokenizer, input_text, generation_config, stopping_criteria_list, device, chunk_delay, max_length=2048):
     encoded_input = tokenizer(input_text, return_tensors="pt", truncation=True, max_length=max_length).to(device)
-    input_length = encoded_input["input_ids"].shape[1]
-    remaining_tokens = max_length - input_length
-
-    if remaining_tokens <= 0:
-        yield ""
-
-    generation_config.max_new_tokens = min(remaining_tokens, generation_config.max_new_tokens)
-
-    def stop_criteria(input_ids, scores):
-        decoded_output = tokenizer.decode(int(input_ids[0][-1]), skip_special_tokens=True)
-        return decoded_output in stop_sequences
-
-    stopping_criteria = StoppingCriteriaList([stop_criteria])
-
-    output_text = ""
-    outputs = model.generate(
-        **encoded_input,
-        do_sample=generation_config.do_sample,
-        max_new_tokens=generation_config.max_new_tokens,
-        temperature=generation_config.temperature,
-        top_p=generation_config.top_p,
-        top_k=generation_config.top_k,
-        repetition_penalty=generation_config.repetition_penalty,
-        num_return_sequences=generation_config.num_return_sequences,
-        stopping_criteria=stopping_criteria,
-        output_scores=True,
-        return_dict_in_generate=True
-    )
-
-    for output in outputs.sequences:
-        for token_id in output:
-            token = tokenizer.decode(token_id, skip_special_tokens=True)
-            yield token
-            await asyncio.sleep(chunk_delay)
-
-        if stop_sequences and any(stop in output_text for stop in stop_sequences):
-            yield output_text
-            return
-
-        outputs = model.generate(
+
+    with torch.no_grad():
+        streamer = model.generate(
             **encoded_input,
-            do_sample=generation_config.do_sample,
-            max_new_tokens=generation_config.max_new_tokens,
-            temperature=generation_config.temperature,
-            top_p=generation_config.top_p,
-            top_k=generation_config.top_k,
-            repetition_penalty=generation_config.repetition_penalty,
-            num_return_sequences=generation_config.num_return_sequences,
-            stopping_criteria=stopping_criteria,
-            output_scores=True,
-            return_dict_in_generate=True
+            generation_config=generation_config,
+            stopping_criteria=stopping_criteria_list,
+            stream=True, # Ensure streaming is enabled if supported by the model
+            return_dict_in_generate=True,
+            output_scores=True
         )
 
+        for output in streamer.sequences[:, encoded_input["input_ids"].shape[-1]:]: # Stream from the *new* tokens
+            token = tokenizer.decode(output, skip_special_tokens=True)
+            if token: # Avoid yielding empty tokens
+                yield token
+                await asyncio.sleep(chunk_delay)
+
+
 @app.post("/generate-image")
 async def generate_image(request: GenerateRequest):
     try: