Update app.py

app.py

@@ -1,18 +1,11 @@
 import os
 import torch
 from fastapi import FastAPI, HTTPException
-from fastapi.responses import StreamingResponse
+from fastapi.responses import JSONResponse
 from pydantic import BaseModel, field_validator
-from transformers import (
-    AutoConfig,
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    GenerationConfig,
-    StoppingCriteriaList
-)
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, GenerationConfig, StoppingCriteriaList
 import boto3
 import uvicorn
-import asyncio
 from io import BytesIO
 from transformers import pipeline
 
@@ -26,21 +19,26 @@ s3_client = boto3.client('s3', aws_access_key_id=AWS_ACCESS_KEY_ID, aws_secret_a
 
 app = FastAPI()
 
+SPECIAL_TOKENS = {
+    "bos_token": "<|startoftext|>",
+    "eos_token": "<|endoftext|>",
+    "pad_token": "[PAD]",
+    "unk_token": "[UNK]",
+}
+
 class GenerateRequest(BaseModel):
     model_name: str
     input_text: str = ""
     task_type: str
     temperature: float = 1.0
     max_new_tokens: int = 10
-    stream: bool = True
     top_p: float = 1.0
     top_k: int = 50
-    repetition_penalty: float = 1.1  # Increased default to discourage repetition
+    repetition_penalty: float = 1.1
     num_return_sequences: int = 1
     do_sample: bool = True
-    chunk_delay: float = 0.0
     stop_sequences: list[str] = []
-    no_repeat_ngram_size: int = 2 # Add parameter to prevent repeating ngrams
+    no_repeat_ngram_size: int = 2
 
     @field_validator("model_name")
     def model_name_cannot_be_empty(cls, v):
@@ -62,11 +60,11 @@ class GenerateRequest(BaseModel):
         return v
 
 class S3ModelLoader:
-    def __init__(self, bucket_name, s3_client):
+    def.__init__(self, bucket_name, s3_client):
         self.bucket_name = bucket_name
         self.s3_client = s3_client
 
-    def _get_s3_uri(self, model_name):
+    def._get_s3_uri(self, model_name):
         return f"s3://{self.bucket_name}/{model_name.replace('/', '-')}"
 
     async def load_model_and_tokenizer(self, model_name):
@@ -75,20 +73,20 @@ class S3ModelLoader:
             config = AutoConfig.from_pretrained(s3_uri, local_files_only=True)
             model = AutoModelForCausalLM.from_pretrained(s3_uri, config=config, local_files_only=True)
             tokenizer = AutoTokenizer.from_pretrained(s3_uri, config=config, local_files_only=True)
-
-            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
-
+            tokenizer.add_special_tokens(SPECIAL_TOKENS)
+            model.resize_token_embeddings(len(tokenizer))
+            if tokenizer.pad_token_id is None:
+                tokenizer.pad_token_id = tokenizer.eos_token_id
             return model, tokenizer
         except EnvironmentError:
             try:
                 config = AutoConfig.from_pretrained(model_name)
                 tokenizer = AutoTokenizer.from_pretrained(model_name, config=config)
+                tokenizer.add_special_tokens(SPECIAL_TOKENS)
                 model = AutoModelForCausalLM.from_pretrained(model_name, config=config)
-
-                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
-
+                model.resize_token_embeddings(len(tokenizer))
+                if tokenizer.pad_token_id is None:
+                    tokenizer.pad_token_id = tokenizer.eos_token_id
                 model.save_pretrained(s3_uri)
                 tokenizer.save_pretrained(s3_uri)
                 return model, tokenizer
@@ -105,13 +103,11 @@ async def generate(request: GenerateRequest):
         task_type = request.task_type
         temperature = request.temperature
         max_new_tokens = request.max_new_tokens
-        stream = request.stream
         top_p = request.top_p
         top_k = request.top_k
         repetition_penalty = request.repetition_penalty
         num_return_sequences = request.num_return_sequences
         do_sample = request.do_sample
-        chunk_delay = request.chunk_delay
         stop_sequences = request.stop_sequences
         no_repeat_ngram_size = request.no_repeat_ngram_size
 
@@ -127,74 +123,41 @@ async def generate(request: GenerateRequest):
             repetition_penalty=repetition_penalty,
             do_sample=do_sample,
             num_return_sequences=num_return_sequences,
-            no_repeat_ngram_size=no_repeat_ngram_size, # Added to generation config
+            no_repeat_ngram_size=no_repeat_ngram_size,
+            pad_token_id=tokenizer.pad_token_id
         )
 
-        return StreamingResponse(
-            stream_text(model, tokenizer, input_text, generation_config, stop_sequences, device, chunk_delay),
-            media_type="text/plain"
-        )
+        generated_text = generate_text(model, tokenizer, input_text, generation_config, stop_sequences, device)
+        return JSONResponse({"text": generated_text})
 
     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):
+def generate_text(model, tokenizer, input_text, generation_config, stop_sequences, device):
     max_model_length = model.config.max_position_embeddings
     encoded_input = tokenizer(input_text, return_tensors="pt", max_length=max_model_length, truncation=True).to(device)
 
-    def stop_criteria(input_ids, scores):
-        decoded_output = tokenizer.decode(input_ids[0], skip_special_tokens=True)
-        for stop in stop_sequences:
-            if decoded_output.endswith(stop):
-                return True
-        return False
+    stopping_criteria = StoppingCriteriaList()
 
-    stopping_criteria = StoppingCriteriaList([stop_criteria])
+    class CustomStoppingCriteria(StoppingCriteriaList):
+        def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+            decoded_output = tokenizer.decode(input_ids[0], skip_special_tokens=True)
+            for stop in stop_sequences:
+                if decoded_output.endswith(stop):
+                    return True
+            return False
 
-    token_buffer = []
-    output_ids = encoded_input.input_ids
-    while True:
-        try:
-            outputs = model.generate(
-                output_ids,
-                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,
-                pad_token_id=tokenizer.pad_token_id,
-                no_repeat_ngram_size=generation_config.no_repeat_ngram_size, # Passed to model.generate
-            )
-        except IndexError as e:
-            print(f"IndexError during generation: {e}")
-            break
-
-        new_token_ids = outputs.sequences[0][encoded_input.input_ids.shape[-1]:]
-
-        for token_id in new_token_ids:
-            token = tokenizer.decode(token_id, skip_special_tokens=True)
-            token_buffer.append(token)
-            if len(token_buffer) >= 10:
-                yield "".join(token_buffer)
-                token_buffer = []
-            await asyncio.sleep(chunk_delay)
-
-        if token_buffer:
-            yield "".join(token_buffer)
-            token_buffer = []
-
-        if stop_criteria(outputs.sequences, None):
-            break
-
-        if len(new_token_ids) < generation_config.max_new_tokens:
-            break
-
-        output_ids = outputs.sequences
+    stopping_criteria.append(CustomStoppingCriteria())
+
+    outputs = model.generate(
+        encoded_input.input_ids,
+        generation_config=generation_config,
+        stopping_criteria=stopping_criteria,
+        pad_token_id=generation_config.pad_token_id
+    )
+
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return generated_text
 
 @app.post("/generate-image")
 async def generate_image(request: GenerateRequest):