Delete app.py

app.py

@@ -1,1570 +0,0 @@
-import os
-import gc
-import json
-import random
-import torch
-import asyncio
-import logging
-import time
-from typing import List, Dict, Any, Optional, Union, AsyncGenerator, Tuple
-from fastapi import FastAPI, HTTPException, Query, Request, Depends, status
-from fastapi.responses import StreamingResponse, PlainTextResponse, HTMLResponse, JSONResponse
-from fastapi.security import APIKeyHeader
-from pydantic import BaseModel, Field, ValidationError, validator
-from transformers import (
-    AutoConfig, AutoModelForCausalLM, AutoTokenizer,
-    GenerationConfig, LogitsProcessorList,
-    MinLengthLogitsProcessor, MaxLengthCriteria,
-    StoppingCriteriaList, StoppingCriteria
-)
-import uvicorn
-from concurrent.futures import ThreadPoolExecutor
-import math
-import torch.nn.functional as F
-import copy
-
-
-class StopSequenceCriteria(StoppingCriteria):
-    def __init__(self, stop_sequences: List[str], tokenizer: AutoTokenizer):
-        self.tokenizer = tokenizer
-        self.stop_sequences_text = []
-        self.stop_sequence_ids = []
-        for seq in stop_sequences:
-            if seq:
-                encoded_ids = tokenizer.encode(seq, add_special_tokens=False)
-                decoded_text = tokenizer.decode(encoded_ids, skip_special_tokens=True)
-                if decoded_text:
-                    self.stop_sequences_text.append(decoded_text)
-                    self.stop_sequence_ids.append(encoded_ids)
-
-
-    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
-        if not self.stop_sequence_ids:
-            return False
-
-        input_ids_list = input_ids[0].tolist()
-
-        for stop_seq_ids in self.stop_sequence_ids:
-            stop_len = len(stop_seq_ids)
-            if len(input_ids_list) >= stop_len:
-                if input_ids_list[-stop_len:] == stop_seq_ids:
-                    return True
-
-        check_tail_len = max(len(self.tokenizer.decode(seq, add_special_tokens=False)) for seq in self.stop_sequence_ids) + 5 if self.stop_sequence_ids else 50
-        tail_ids = input_ids_list[-min(check_tail_len, len(input_ids_list)):]
-        tail_text = self.tokenizer.decode(tail_ids, skip_special_tokens=True)
-
-        for stop_seq_text in self.stop_sequences_text:
-            if stop_seq_text and stop_seq_text in tail_text:
-                return True
-
-        return False
-
-
-logging.getLogger("uvicorn").handlers.clear()
-logging.getLogger("uvicorn.error").handlers.clear()
-logging.getLogger("uvicorn.access").handlers.clear()
-logging.getLogger("uvicorn").propagate = False
-logging.getLogger("uvicorn.error").propagate = False
-logging.getLogger("uvicorn.access").propagate = False
-logging.getLogger("uvicorn").setLevel(logging.CRITICAL)
-logging.getLogger("uvicorn.error").setLevel(logging.CRITICAL)
-logging.getLogger("uvicorn.access").setLevel(logging.CRITICAL)
-logging.getLogger("fastapi").setLevel(logging.CRITICAL)
-logging.getLogger("transformers").setLevel(logging.CRITICAL)
-logging.getLogger().handlers.clear()
-logging.getLogger().addHandler(logging.NullHandler())
-
-
-MODEL_NAME = os.environ.get("MODEL_NAME", "jnjj/multi-dataset-model")
-SYSTEM_PROMPT = os.environ.get("SYSTEM_PROMPT", "Eres un asistente profesional y servicial.")
-MAX_CONTEXT_TOKENS = int(os.environ.get("MAX_CONTEXT_TOKENS", 1024))
-MAX_GENERATION_TOKENS = int(os.environ.get("MAX_GENERATION_TOKENS", 512))
-MAX_CONCURRENT_GENERATIONS = int(os.environ.get("MAX_CONCURRENT_GENERATIONS", 4))
-TRUST_REMOTE_CODE = os.environ.get("TRUST_REMOTE_CODE", "false").lower() == "true"
-ENABLE_FLASH_ATTENTION_2 = os.environ.get("ENABLE_FLASH_ATTENTION_2", "false").lower() == "true"
-TORCH_DTYPE_STR = os.environ.get("TORCH_DTYPE", "float16")
-TORCH_DTYPE = getattr(torch, TORCH_DTYPE_STR.lower(), torch.float16)
-if TORCH_DTYPE_STR.lower() == "bfloat16" and (not torch.cuda.is_available() or (torch.cuda.is_available() and not torch.cuda.get_device_properties(0).has_bfloat16)):
-     TORCH_DTYPE = torch.float16
-elif not isinstance(TORCH_DTYPE, torch.dtype):
-     TORCH_DTYPE = torch.float16
-
-API_KEY = os.environ.get("API_KEY")
-
-global_model = None
-global_tokenizer = None
-global_tokens: Dict[str, Optional[int]] = {}
-executor = ThreadPoolExecutor(max_workers=MAX_CONCURRENT_GENERATIONS)
-generation_semaphore = asyncio.Semaphore(MAX_CONCURRENT_GENERATIONS)
-
-api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
-
-async def get_api_key(api_key: str = Depends(api_key_header)):
-    if API_KEY is None:
-        return
-    if api_key is None or api_key != API_KEY:
-        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid or missing API Key")
-    return api_key
-
-
-class GenerateRequest(BaseModel):
-    input_text: str = Field(..., description="The input text from the user.")
-    history: Optional[List[Dict[str, str]]] = Field(None, description="Conversation history in the format [{\"role\": \"system\", \"content\": \"...\"}, {\"role\": \"user\", \"content\": \"...\"}, {\"role\": \"assistant\", \"content\": \"...\"}, ...]. The last message in history should be the user's most recent previous turn if provided.")
-    stream: bool = Field(True, description="Whether to stream the response.")
-    temperature: float = Field(1.0, ge=0.0, le=2.0, description="Controls the randomness of the output. Higher values mean more random output.")
-    top_k: int = Field(50, ge=0, description="The number of highest probability vocabulary tokens to keep for top-k-filtering.")
-    top_p: float = Field(1.0, ge=0.0, le=1.0, description="If set to float < 1, only the most probable tokens with probabilities that add up to top_p or higher are kept for generation.")
-    repetition_penalty: float = Field(1.0, ge=0.0, description="The parameter for repetition penalty. 1.0 means no penalty.")
-    frequency_penalty: float = Field(0.0, ge=0.0, description="Penalizes tokens based on their frequency in the generation.")
-    presence_penalty: float = Field(0.0, ge=0.0, description="Penalizes tokens based on whether they appear in the generation.")
-    num_beams: int = Field(1, ge=1, description="Number of beams for beam search. 1 means no beam search.")
-    length_penalty: float = Field(1.0, ge=0.0, description="Exponential penalty to the length. 1.0 means no penalty. Set to values greater than 1.0 in order to encourage longer sequences, and to values less than 1.0 in order to encourage shorter sequences.")
-    no_repeat_ngram_size: int = Field(0, ge=0, description="If set to int > 0, all ngrams of that size can only occur once.")
-    early_stopping: bool = Field(False, description="Whether to stop the beam search when at least one beam reaches a stopping criteria (e.g. EOS token).")
-    do_sample: bool = Field(True, description="Whether to use sampling.")
-    use_mirostat: bool = Field(False, description="Whether to use Mirostat sampling.")
-    mirostat_tau: float = Field(5.0, ge=0.0, description="Mirostat tau parameter.")
-    mirostat_eta: float = Field(0.1, ge=0.0, description="Mirostat eta parameter.")
-    max_new_tokens: int = Field(MAX_GENERATION_TOKENS, ge=1, description="The maximum number of tokens to generate.")
-    system_prompt: Optional[str] = Field(None, description="Override the default system prompt for this request.")
-    seed: Optional[int] = Field(None, description="Random seed for generation.")
-    stop_sequences: Optional[List[str]] = Field(None, description="List of strings that will stop generation.")
-    tokenize_only: bool = Field(False, description="If true, only tokenize the input and return token count, not generate.")
-    strip_trailing_whitespace: bool = Field(False, description="Whether to strip trailing whitespace from the final generated text.")
-    remove_incomplete_sentences: bool = Field(False, description="Whether to remove the last sentence if it appears incomplete.")
-    num_return_sequences: int = Field(1, ge=1, le=5, description="The number of sequences to return for non-streaming generation (ignored for streaming).")
-    bad_words_ids: Optional[List[List[int]]] = Field(None, description="List of token ids that are not allowed to be generated.")
-    forced_bos_token_id: Optional[int] = Field(None, description="The id of the token to force as the first generated token.")
-    forced_eos_token_id: Optional[int] = Field(None, description="The id of the token to force as the last generated token.")
-    renormalize_logits: Optional[bool] = Field(None, description="Whether to renormalize the logits before applying penalties and sampling.")
-    suppress_tokens: Optional[List[int]] = Field(None, description="List of tokens that will be suppressed at generation.")
-    begin_suppress_tokens: Optional[List[int]] = Field(None, description="List of tokens that will be suppressed at the beginning of the generation.")
-    end_suppress_tokens: Optional[List[int]] = Field(None, description="List of tokens that will be suppressed at the end of the generation.")
-    encoder_no_repeat_ngram_size: int = Field(0, ge=0, description="If set to int > 0, all ngrams of that size can only occur once in the encoder_output.")
-    min_length: int = Field(0, ge=0, description="The minimum length of the sequence to be generated.")
-    max_length: Optional[int] = Field(None, description="The maximum total length (input + generated) of the sequence to be generated.")
-    exponential_decay_length_penalty: Optional[Tuple[float, int, float]] = Field(None, description="Exponential decay length penalty parameters (decay_factor, start_index, threshold).")
-    use_cache: bool = Field(True, description="Whether or not the model should use the past key/values attentions to speed up decoding.")
-    top_k_dynamic: Optional[int] = Field(None, description="Dynamic top_k parameter.")
-    top_p_dynamic: Optional[float] = Field(None, description="Dynamic top_p parameter.")
-    typical_p: float = Field(1.0, ge=0.0, le=1.0, description="Typical P sampling parameter.")
-    epsilon_cutoff: float = Field(0.0, ge=0.0, description="Epsilon cutoff parameter for Locally Typical Sampling.")
-    eta_cutoff: float = Field(0.0, ge=0.0, description="Eta cutoff parameter for Locally Typical Sampling.")
-    temperature_cutoff: Optional[float] = Field(None, ge=0.0, description="Temperature cutoff parameter.")
-    encoder_repetition_penalty: float = Field(1.0, ge=0.0, description="The parameter for repetition penalty on encoder_outputs.")
-    max_time: Optional[float] = Field(None, ge=0.0, description="The maximum amount of time in seconds the generation process can take.")
-    output_watermark: bool = Field(False, description="Whether or not to output watermark information (if supported by the model).")
-    remove_input_from_output: bool = Field(False, description="Whether to remove the input tokens from the output.")
-    eos_token_id_override: Optional[int] = Field(None, description="Override the EOS token id for generation.")
-    pad_token_id_override: Optional[int] = Field(None, description="Override the PAD token id for generation.")
-    bos_token_id_override: Optional[int] = Field(None, description="Override the BOS token id for generation.")
-    repetition_penalty_range: Optional[int] = Field(None, ge=0, description="The number of tokens to consider for repetition penalty.")
-    diversity_penalty: float = Field(0.0, ge=0.0, description="The diversity penalty parameter for diverse beam search.")
-    num_beam_groups: int = Field(1, ge=1, description="Number of groups to divide beams into in diverse beam search.")
-
-    return_dict_in_generate: bool = Field(False, description="Whether or not to return a dictionary from model.generate().")
-    output_attentions: bool = Field(False, description="Whether or not to return attentions from model.generate().")
-    output_hidden_states: bool = Field(False, description="Whether or not to return hidden states from model.generate().")
-    output_scores: bool = Field(False, description="Whether or not to return scores from model.generate().")
-    return_token_logprobs: bool = Field(False, description="Whether to return token logprobs in the streaming response (only for stream=True).")
-    return_text_from_sequence: bool = Field(True, description="Whether to decode the sequence back into text in the response.")
-    length_normalization_factor: Optional[float] = Field(None, description="The length normalization factor for beam search.")
-    min_new_tokens: int = Field(0, ge=0, description="The minimum number of new tokens to generate.")
-    do_normalize_logits: bool = Field(False, description="Whether to normalize the logits before applying sampling/greedy decoding.")
-
-    return_generation_inputs: bool = Field(False, description="Whether or not to return generation inputs from model.generate().")
-    return_unused_generate_parameters: bool = Field(False, description="Whether or not to return unused generate parameters from model.generate().")
-    use_fast_tokenizer: bool = Field(True, description="Whether to use a fast tokenizer if available.")
-    model_kwargs: Optional[Dict[str, Any]] = Field(None, description="Additional model keyword arguments for generate().")
-    tokenizer_kwargs: Optional[Dict[str, Any]] = Field(None, description="Additional tokenizer keyword arguments for encode().")
-
-    return_only_text: bool = Field(False, description="If true, only return the generated text (plain text for streaming, JSON list of strings for non-streaming num_return_sequences > 1). Other metadata will be omitted.")
-
-
-    @validator('stop_sequences')
-    def validate_stop_sequences(cls, v):
-        if v is not None:
-            if not all(isinstance(seq, str) for seq in v):
-                raise ValueError('Each stop sequence must be a string')
-        return v
-
-    @validator('bad_words_ids')
-    def validate_bad_words_ids(cls, v):
-        if v is not None:
-            if not all(isinstance(word_id_list, list) and all(isinstance(token_id, int) for token_id in word_id_list) for word_id_list in v):
-                raise ValueError('bad_words_ids must be a list of lists of integers')
-        return v
-
-    @validator('exponential_decay_length_penalty')
-    def validate_exponential_decay_length_penalty(cls, v):
-        if v is not None:
-            if not (isinstance(v, (list, tuple)) and len(v) == 3 and
-                    isinstance(v[0], (int, float)) and v[0] > 0 and
-                    isinstance(v[1], int) and v[1] >= 0 and
-                    isinstance(v[2], (int, float))):
-                 raise ValueError('exponential_decay_length_penalty must be a tuple/list of 3 numbers (decay_factor, start_index, threshold)')
-        return v
-
-
-class TokenizeRequest(BaseModel):
-    text: Union[str, List[str]] = Field(..., description="Text or list of texts to tokenize.")
-    add_special_tokens: bool = Field(True, description="Whether to add special tokens.")
-    is_split_into_words: bool = Field(False, description="Whether the input text is pre-tokenized.")
-    return_token_type_ids: bool = Field(False, description="Whether to return token type IDs.")
-    padding: Union[bool, str] = Field(False, description="Enable padding. Can be 'longest', 'max_length', 'do_not_pad' or True/False.")
-    truncation: Union[bool, str] = Field(False, description="Enable truncation. Can be 'only_first', 'only_second', 'longest_first', 'do_not_truncate' or True/False.")
-    max_length: Optional[int] = Field(None, ge=1, description="Maximum length for padding and truncation.")
-    return_tensors: Optional[str] = Field(None, description="The type of tensors to return ('pt' for PyTorch, 'tf' for TensorFlow, 'np' for NumPy).")
-    return_attention_mask: Optional[bool] = Field(None, description="Whether to return the attention mask.")
-    return_offsets_mapping: Optional[bool] = Field(None, description="Whether to return offsets mapping.")
-    return_length: Optional[bool] = Field(None, description="Whether to return the length of the encoded sequences.")
-    verbose: bool = Field(False, description="Whether to print verbose information.")
-    tokenizer_kwargs: Optional[Dict[str, Any]] = Field(None, description="Additional tokenizer keyword arguments.")
-
-
-class DecodeRequest(BaseModel):
-    token_ids: List[int] = Field(..., description="List of token IDs to decode.")
-    skip_special_tokens: bool = Field(True, description="Whether to skip special tokens during decoding.")
-    clean_up_tokenization_spaces: bool = Field(True, description="Whether to remove leading/trailing whitespace and double spaces.")
-    decode_kwargs: Optional[Dict[str, Any]] = Field(None, description="Additional decode keyword arguments.")
-
-
-class SystemPromptUpdateRequest(BaseModel):
-    system_prompt: str = Field(..., description="The new global system prompt.")
-
-class ModelReloadRequest(BaseModel):
-    model_name: Optional[str] = Field(None, description="New model name to load. If None, reloads the current model.")
-    trust_remote_code: Optional[bool] = Field(None, description="Override trust_remote_code setting.")
-    enable_flash_attention_2: Optional[bool] = Field(None, description="Override enable_flash_attention_2 setting.")
-    torch_dtype: Optional[str] = Field(None, description="Override torch_dtype setting (e.g., 'float16', 'bfloat16').")
-    model_kwargs: Optional[Dict[str, Any]] = Field(None, description="Additional model keyword arguments for from_pretrained().")
-    tokenizer_kwargs: Optional[Dict[str, Any]] = Field(None, description="Additional tokenizer keyword arguments for from_pretrained().")
-
-def format_conversation(input_text: str, history: Optional[List[Dict[str, str]]], system_prompt: Optional[str]) -> str:
-    full_history: List[Dict[str, str]] = []
-    used_system_prompt = system_prompt if system_prompt is not None else SYSTEM_PROMPT
-    if not history or history[0].get("role") != "system" or history[0].get("content") != used_system_prompt:
-         full_history.append({"role": "system", "content": used_system_prompt})
-
-    if history:
-        full_history.extend(history)
-
-    if not full_history or full_history[-1].get("role") != "user" or full_history[-1].get("content") != input_text:
-        full_history.append({"role": "user", "content": input_text})
-
-    if global_tokenizer and hasattr(global_tokenizer, 'apply_chat_template') and global_tokenizer.chat_template:
-         try:
-              return global_tokenizer.apply_chat_template(full_history, tokenize=False, add_generation_prompt=True)
-         except Exception as e:
-              logging.error(f"Failed to apply chat template: {e}. Falling back to manual formatting.")
-              pass
-
-    formatted_text = ""
-    for i, message in enumerate(full_history):
-        if i == 0 and message["role"] == "system" and len(full_history) > 1 and full_history[1].get("role") == "system":
-             continue
-
-        if message["role"] == "system":
-            formatted_text += f"{message['content'].strip()}\n\n"
-        elif message["role"] == "user":
-            formatted_text += f"Usuario: {message['content'].strip()}\n"
-        elif message["role"] == "assistant":
-             formatted_text += f"Bot: {message['content'].strip()}\n"
-
-    if not formatted_text.endswith("Bot:"):
-         formatted_text += "Bot:"
-
-    return formatted_text.strip()
-
-def truncate_encoded_ids(input_ids: torch.Tensor, max_length: int) -> torch.Tensor:
-    if input_ids.shape[-1] > max_length:
-        return input_ids[:, -max_length:]
-    return input_ids
-
-def apply_seed(seed: Optional[int]):
-    if seed is not None:
-        torch.manual_seed(seed)
-        if torch.cuda.is_available():
-            torch.cuda.manual_seed_all(seed)
-        random.seed(seed)
-
-def get_stopping_criteria(req: GenerateRequest, initial_ids: torch.Tensor, tokenizer: AutoTokenizer) -> StoppingCriteriaList:
-    criteria = StoppingCriteriaList()
-    max_len_from_req = None
-    if req.max_length is not None and req.max_length > 0:
-         max_len_from_req = req.max_length
-    elif req.max_new_tokens is not None and req.max_new_tokens > 0:
-         max_len_from_req = initial_ids.shape[-1] + req.max_new_tokens
-    else:
-         max_len_from_req = initial_ids.shape[-1] + MAX_GENERATION_TOKENS
-
-
-    if max_len_from_req is not None and max_len_from_req > 0:
-         criteria.append(MaxLengthCriteria(max_len_from_req))
-
-
-    if req.min_length is not None and req.min_length > 0:
-        eos_token_id = req.eos_token_id_override if req.eos_token_id_override is not None else global_tokens.get("eos_token_id", -1)
-        criteria.append(MinLengthLogitsProcessor(initial_ids.shape[-1] + req.min_length, eos_token_id))
-
-    if req.stop_sequences:
-        criteria.append(StopSequenceCriteria(req.stop_sequences, tokenizer))
-
-
-    return criteria
-
-
-def generate_next_token_sync(
-    input_ids,
-    past_key_values,
-    gen_cfg: GenerationConfig,
-    device: str
-) -> Tuple[torch.Tensor, Any, Optional[float], Optional[torch.Tensor], Any, Any]:
-    with torch.no_grad():
-        outputs = global_model(
-            input_ids, past_key_values=past_key_values,
-            use_cache=gen_cfg.use_cache, return_dict=True,
-            output_attentions=gen_cfg.output_attentions,
-            output_hidden_states=gen_cfg.output_hidden_states,
-            output_scores=gen_cfg.output_scores,
-        )
-        logits = outputs.logits[:, -1, :]
-        past = outputs.past_key_values
-
-        scores = outputs.scores if gen_cfg.output_scores else None
-        attentions = outputs.attentions if gen_cfg.output_attentions else None
-        hidden_states = outputs.hidden_states if gen_cfg.output_hidden_states else None
-
-        step_logits_for_criteria = logits.clone()
-
-        if gen_cfg.do_normalize_logits:
-             logits = F.log_softmax(logits, dim=-1)
-
-        if gen_cfg.do_sample:
-            if gen_cfg.use_mirostat_mode == 1 and hasattr(global_model, 'mirostat_sample_logits'):
-                 token = global_model.mirostat_sample_logits(
-                      logits=logits,
-                      temperature=gen_cfg.temperature,
-                      mirostat_tau=gen_cfg.mirostat_tau,
-                      mirostat_eta=gen_cfg.mirostat_eta
-                 ).unsqueeze(0).to(device)
-            else:
-                 logits = logits / gen_cfg.temperature
-
-                 if gen_cfg.temperature_cutoff is not None and gen_cfg.temperature_cutoff > 0:
-                      logits = torch.where(logits < gen_cfg.temperature_cutoff, torch.tensor(-float('Inf')).to(logits.device), logits)
-
-                 if gen_cfg.top_k:
-                     topk_values, topk_indices = torch.topk(logits, gen_cfg.top_k)
-                     logits[logits < topk_values[:, -1]] = -float('Inf')
-
-                 if gen_cfg.top_p < 1.0:
-                      sorted_logits, sorted_indices = torch.sort(logits, dim=-1, descending=True)
-                      cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
-                      sorted_indices_to_remove = cumulative_probs > gen_cfg.top_p
-                      sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
-                      sorted_indices_to_remove[..., 0] = False
-                      indices_to_remove = sorted_indices[sorted_indices_to_remove]
-                      logits[:, indices_to_remove] = -float('Inf')
-
-                 if gen_cfg.typical_p < 1.0:
-                     probs = torch.softmax(logits, dim=-1)
-                     entropy = torch.distributions.Categorical(probs).entropy()
-                     probs_sorted, indices_sorted = torch.sort(probs, dim=-1, descending=True)
-                     cumsum_probs_sorted = torch.cumsum(probs_sorted, dim=-1)
-                     mask = cumsum_probs_sorted < gen_cfg.typical_p * entropy.exp()
-                     indices_to_remove = indices_sorted[~mask]
-                     logits[:, indices_to_remove] = -float('Inf')
-
-
-                 if gen_cfg.epsilon_cutoff is not None and gen_cfg.epsilon_cutoff > 0:
-                      probs = torch.softmax(logits, dim=-1)
-                      mask = probs < gen_cfg.epsilon_cutoff
-                      logits[:, mask] = -float('Inf')
-
-                 if gen_cfg.eta_cutoff is not None and gen_cfg.eta_cutoff > 0:
-                      probs = torch.softmax(logits, dim=-1)
-                      mask = probs > gen_cfg.eta_cutoff
-                      logits[:, ~mask] = -float('Inf')
-
-                 probs = torch.softmax(logits, dim=-1)
-                 token = torch.multinomial(probs, 1)
-        else:
-            token = torch.argmax(logits, dim=-1, keepdim=True)
-
-        token_logprob = None
-        if gen_cfg.output_scores:
-             log_probs = F.log_softmax(step_logits_for_criteria, dim=-1)
-             if 0 <= token.squeeze().item() < log_probs.shape[-1]:
-                 token_logprob = float(log_probs[:, token.squeeze()].item())
-             else:
-                  token_logprob = None
-
-
-    return token, past, token_logprob, step_logits_for_criteria, attentions, hidden_states
-
-def post_process_text(text: str, strip_trailing_whitespace: bool, remove_incomplete_sentences: bool) -> str:
-    if strip_trailing_whitespace:
-        text = text.rstrip()
-    if remove_incomplete_sentences:
-        for terminator in ['.', '!', '?', '\n']:
-            last_terminator = text.rfind(terminator)
-            if last_terminator != -1:
-                text = text[:last_terminator + 1]
-                break
-    return text
-
-async def stream_generation_logic(req: GenerateRequest, initial_ids: torch.Tensor, gen_cfg: GenerationConfig, device: str) -> AsyncGenerator[str, None]:
-    past = None
-    generated_tokens_count = 0
-    eos_token_id = req.eos_token_id_override if req.eos_token_id_override is not None else global_tokens.get("eos_token_id")
-    pad_token_id = req.pad_token_id_override if req.pad_token_id_override is not None else global_tokens.get("pad_token_id", eos_token_id)
-    stop_token_ids = {eos_token_id} if eos_token_id is not None else set()
-    if pad_token_id is not None and pad_token_id != eos_token_id:
-        stop_token_ids.add(pad_token_id)
-
-    current_ids = initial_ids
-    start_time = time.time()
-    total_ids_list = initial_ids.tolist()[0]
-    finish_reason = "unknown"
-
-    stopping_criteria = get_stopping_criteria(req, initial_ids, global_tokenizer)
-
-    last_step_logits = None
-    accumulated_text_for_processing = ""
-
-    try:
-        while True:
-            if generated_tokens_count >= req.max_new_tokens:
-                 finish_reason = "max_new_tokens"
-                 break
-            if req.max_time is not None and (time.time() - start_time) > req.max_time:
-                 finish_reason = "time"
-                 break
-
-            input_ids_sync = current_ids if past is None else token
-
-            token, past, token_logprob, step_logits, attentions, hidden_states = await asyncio.to_thread(
-                generate_next_token_sync,
-                input_ids_sync,
-                past,
-                gen_cfg,
-                device
-            )
-            last_step_logits = step_logits
-
-            generated_token_id = token[0].item()
-            total_ids_list.append(generated_token_id)
-
-            text = global_tokenizer.decode([generated_token_id], skip_special_tokens=True)
-            accumulated_text_for_processing += text
-
-            chunk_payload: Dict[str, Any] = {
-                "type": "token",
-                "text": text,
-                "token_id": generated_token_id,
-                "generated_tokens_count": generated_tokens_count + 1,
-            }
-            if req.return_token_logprobs and token_logprob is not None:
-                chunk_payload["logprob"] = token_logprob
-
-            yield json.dumps(chunk_payload) + "\n"
-
-            if generated_token_id in stop_token_ids:
-                finish_reason = "eos_token"
-                break
-
-            current_full_ids_tensor = torch.tensor([total_ids_list], device=device)
-            if stopping_criteria(current_full_ids_tensor, step_logits):
-                 finish_reason = "stopping_criteria"
-                 current_len = len(total_ids_list)
-                 initial_len = initial_ids.shape[-1]
-
-                 max_len_crit_met = any(isinstance(c, MaxLengthCriteria) for c in stopping_criteria) and \
-                                     ( (req.max_new_tokens is not None and current_len >= (initial_len + req.max_new_tokens)) or
-                                       (req.max_length is not None and current_len >= req.max_length) )
-                 stop_seq_crit_met = any(isinstance(c, StopSequenceCriteria) for c in stopping_criteria) and req.stop_sequences and \
-                                      any(seq in global_tokenizer.decode(total_ids_list[initial_len:], skip_special_tokens=True) for seq in req.stop_sequences)
-
-                 if max_len_crit_met:
-                      if req.max_new_tokens is not None and current_len >= (initial_len + req.max_new_tokens):
-                          finish_reason = "max_new_tokens"
-                      elif req.max_length is not None and current_len >= req.max_length:
-                          finish_reason = "max_length"
-
-                 if stop_seq_crit_met:
-                      finish_reason = "stop_sequence"
-
-
-                 break
-
-
-            current_ids = token
-            generated_tokens_count += 1
-
-        final_text_raw = global_tokenizer.decode(total_ids_list[initial_ids.shape[-1]:], skip_special_tokens=True)
-        if req.stop_sequences and finish_reason == "stop_sequence":
-             for stop_seq in req.stop_sequences:
-                  if stop_seq and stop_seq in final_text_raw:
-                       final_text_raw = final_text_raw.split(stop_seq, 1)[0]
-                       break
-
-        final_text_processed = post_process_text(final_text_raw, req.strip_trailing_whitespace, req.remove_incomplete_sentences)
-
-
-        final_payload: Dict[str, Any] = {
-            "type": "done",
-            "total_prompt_tokens": initial_ids.shape[-1],
-            "total_generated_tokens": generated_tokens_count,
-            "total_sequence_tokens": len(total_ids_list),
-            "final_text": final_text_processed,
-            "finish_reason": finish_reason
-        }
-        yield json.dumps(final_payload) + "\n"
-
-
-    except Exception as e:
-         logging.exception("Streaming generation error:")
-         error_payload = {"type": "error", "message": str(e)}
-         yield json.dumps(error_payload) + "\n"
-
-    finally:
-        await cleanup(device)
-
-async def plain_text_stream_processor(json_stream_generator: AsyncGenerator[str, None], req: GenerateRequest) -> AsyncGenerator[str, None]:
-    accumulated_text = ""
-    try:
-        async for chunk in json_stream_generator:
-            try:
-                data = json.loads(chunk)
-                if data.get("type") == "token":
-                     yield data.get("text", "")
-                     accumulated_text += data.get("text", "")
-                elif data.get("type") == "done":
-                     final_text_raw = data.get("final_text", "")
-                     return
-                elif data.get("type") == "error":
-                     raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=data.get("message", "An unknown error occurred during streaming."))
-            except json.JSONDecodeError:
-                logging.warning(f"Received non-JSON chunk in stream: {chunk.strip()}")
-            except HTTPException as e:
-                 raise e
-            except Exception as e:
-                 logging.exception("Error processing stream chunk for plain text:")
-                 raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error processing stream chunk: {e}")
-    except Exception as e:
-         raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Streaming error: {e}")
-
-
-async def non_stream_generation_logic(req: GenerateRequest, initial_ids: torch.Tensor, gen_cfg: GenerationConfig, device: str) -> Dict[str, Any]:
-    try:
-        logits_processor_list = LogitsProcessorList()
-
-        stopping_criteria_list = get_stopping_criteria(req, initial_ids, global_tokenizer)
-
-
-        with torch.no_grad():
-             out = global_model.generate(
-                input_ids=initial_ids,
-                generation_config=gen_cfg,
-                return_dict_in_generate=True,
-                output_scores=req.output_scores,
-                output_attentions=req.output_attentions,
-                output_hidden_states=req.output_hidden_states,
-                num_return_sequences=req.num_return_sequences,
-                bad_words_ids=req.bad_words_ids,
-                suppress_tokens=req.suppress_tokens,
-                begin_suppress_tokens=req.begin_suppress_tokens,
-                end_suppress_tokens=req.end_suppress_tokens,
-                logits_processor=logits_processor_list if logits_processor_list else None,
-                stopping_criteria=stopping_criteria_list if stopping_criteria_list else None,
-             )
-
-        generated_data = []
-        for i in range(req.num_return_sequences):
-            if i >= len(out.sequences):
-                 break
-
-            sequence = out.sequences[i]
-            start_index = initial_ids.shape[-1]
-            generated_ids_tensor = sequence[start_index:]
-            full_sequence_ids = sequence.tolist()
-
-            text = global_tokenizer.decode(generated_ids_tensor, skip_special_tokens=True)
-
-            if req.stop_sequences:
-                 for stop_seq in req.stop_sequences:
-                      if stop_seq and stop_seq in text:
-                           text = text.split(stop_seq, 1)[0]
-                           break
-
-            text = post_process_text(text, req.strip_trailing_whitespace, req.remove_incomplete_sentences)
-
-            finish_reason = "length"
-            eos_token_id = req.eos_token_id_override if req.eos_token_id_override is not None else global_tokens.get("eos_token_id")
-            if len(generated_ids_tensor) > 0 and generated_ids_tensor[-1] == eos_token_id:
-                 finish_reason = "eos_token"
-            elif len(generated_ids_tensor) >= gen_cfg.max_new_tokens:
-                 finish_reason = "max_new_tokens"
-            elif req.max_length is not None and len(full_sequence_ids) >= req.max_length:
-                 finish_reason = "max_length"
-            elif hasattr(out, 'max_time_exceeded') and out.max_time_exceeded:
-                 finish_reason = "time"
-
-            if req.stop_sequences and finish_reason == "length":
-                 decoded_full_output = global_tokenizer.decode(full_sequence_ids, skip_special_tokens=True)
-                 if any(seq in decoded_full_output for seq in req.stop_sequences):
-                      finish_reason = "stop_sequence"
-
-
-            item_data: Dict[str, Any] = {
-                "text": text if req.return_text_from_sequence else None,
-                "token_ids": generated_ids_tensor.tolist(),
-                "generated_tokens_count": len(generated_ids_tensor),
-                "finish_reason": finish_reason
-            }
-            if not req.remove_input_from_output:
-                 item_data["full_sequence_token_ids"] = full_sequence_ids
-
-            if req.output_scores and hasattr(out, 'scores') and out.scores is not None:
-                 item_data["scores"] = "Scores output needs custom handling (complex structure)."
-
-                 if req.return_token_logprobs:
-                      item_data["token_logprobs"] = "Token logprobs require parsing scores output which is complex for batched/beamed generation."
-
-            if req.output_attentions and hasattr(out, 'attentions') and out.attentions is not None:
-                 item_data["attentions"] = "Attentions output needs custom handling (too large)."
-            if req.output_hidden_states and hasattr(out, 'hidden_states') and out.hidden_states is not None:
-                 item_data["hidden_states"] = "Hidden states output needs custom handling (too large)."
-            if req.output_watermark and hasattr(out, 'watermark') and out.watermark is not None:
-                 item_data["watermark"] = out.watermark[i] if isinstance(out.watermark, list) and len(out.watermark) > i else out.watermark
-
-
-            generated_data.append(item_data)
-
-
-        response_payload: Dict[str, Any] = {
-             "prompt_tokens": initial_ids.shape[-1],
-             "generated_sequences": generated_data,
-        }
-        if req.num_return_sequences == 1 and generated_data:
-             response_payload["total_tokens"] = response_payload["prompt_tokens"] + generated_data[0]["generated_tokens_count"]
-
-        if req.return_dict_in_generate:
-             raw_out_dict = {}
-             for key in out.keys():
-                  if key not in ['sequences', 'scores', 'attentions', 'hidden_states', 'past_key_values', 'watermark', 'sequences_scores']:
-                       value = out[key]
-                       if isinstance(value, torch.Tensor):
-                            raw_out_dict[key] = value.tolist()
-                       else:
-                            raw_out_dict[key] = value
-
-             response_payload["raw_generate_output"] = raw_out_dict
-
-        return response_payload
-
-    except Exception as e:
-         logging.exception("Non-streaming generation error:")
-         raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Generation error: {e}")
-
-async def cleanup(device: str):
-    if device == "cuda":
-        torch.cuda.empty_cache()
-    gc.collect()
-
-app = FastAPI(title="Chatbot Profesional API", version="1.0.0")
-
-@app.on_event("startup")
-async def load_model():
-    global global_model, global_tokenizer, global_tokens, MODEL_NAME, TRUST_REMOTE_CODE, ENABLE_FLASH_ATTENTION_2, TORCH_DTYPE, TORCH_DTYPE_STR
-
-    try:
-         TORCH_DTYPE = getattr(torch, TORCH_DTYPE_STR.lower(), torch.float16)
-         if TORCH_DTYPE_STR.lower() == "bfloat16" and (not torch.cuda.is_available() or (torch.cuda.is_available() and not torch.cuda.get_device_properties(0).has_bfloat16)):
-              logging.warning(f"Requested dtype {TORCH_DTYPE_STR} but not supported or CUDA not available. Falling back to float16.")
-              TORCH_DTYPE = torch.float16
-         elif not isinstance(TORCH_DTYPE, torch.dtype):
-              logging.warning(f"Invalid TORCH_DTYPE specified: {TORCH_DTYPE_STR}. Falling back to float16.")
-              TORCH_DTYPE = torch.float16
-    except AttributeError:
-         logging.warning(f"Invalid TORCH_DTYPE specified: {TORCH_DTYPE_STR}. Falling back to float16.")
-         TORCH_DTYPE = torch.float16
-
-    try:
-        logging.info(f"Loading config for model: {MODEL_NAME}")
-        config = AutoConfig.from_pretrained(MODEL_NAME, trust_remote_code=TRUST_REMOTE_CODE)
-        logging.info("Config loaded.")
-
-        logging.info(f"Loading tokenizer for model: {MODEL_NAME}")
-        tokenizer_kwargs = {"config": config, "trust_remote_code": TRUST_REMOTE_CODE}
-        global_tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, **tokenizer_kwargs)
-        logging.info("Tokenizer loaded.")
-
-        logging.info(f"Loading model: {MODEL_NAME} with dtype {TORCH_DTYPE} (CUDA: {torch.cuda.is_available()})")
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-
-        model_kwargs = {"config": config, "torch_dtype": TORCH_DTYPE, "trust_remote_code": TRUST_REMOTE_CODE}
-        if ENABLE_FLASH_ATTENTION_2 and device == "cuda":
-            try:
-                 if hasattr(config, '_supports_flash_attn_2') and config._supports_flash_attn_2:
-                     model_kwargs['use_flash_attention_2'] = True
-                     logging.info("Flash Attention 2 enabled.")
-                 elif hasattr(AutoModelForCausalLM, '_supports_flash_attn_2') and AutoModelForCausalLM._supports_flash_attn_2(config):
-                      model_kwargs['use_flash_attention_2'] = True
-                      logging.info("Flash Attention 2 enabled via model check.")
-                 else:
-                     logging.warning("Flash Attention 2 requested but not supported by model config.")
-            except Exception as e:
-                 logging.warning(f"Error checking for Flash Attention 2 support: {e}")
-                 pass
-
-
-        global_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, **model_kwargs)
-        global_model.to(device)
-        global_model.eval()
-        logging.info("Model loaded successfully.")
-
-        global_tokens["eos_token_id"] = global_tokenizer.eos_token_id
-        global_tokens["pad_token_id"] = global_tokenizer.pad_token_id
-        if global_tokens["pad_token_id"] is None and global_tokens["eos_token_id"] is not None:
-             global_tokens["pad_token_id"] = global_tokens["eos_token_id"]
-             if global_model.config.pad_token_id is None:
-                  global_model.config.pad_token_id = global_tokens["pad_token_id"]
-        elif global_tokens["pad_token_id"] is None and global_tokens["eos_token_id"] is None:
-             logging.warning("Neither EOS nor PAD token is defined for this tokenizer/model.")
-
-    except Exception as e:
-        logging.exception("Failed to load model or tokenizer:")
-        global_model = None
-        global_tokenizer = None
-        global_tokens = {}
-
-html_code = """
-<!DOCTYPE html>
-<html lang="es">
-<head>
-    <meta charset="UTF-8" />
-    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
-    <title>Chatbot Profesional</title>
-    <style>
-        body { font-family: Arial, sans-serif; margin: 20px; }
-        #chatbox { width: 100%; height: 400px; border: 1px solid #ccc; padding: 10px; overflow-y: scroll; margin-bottom: 10px; }
-        #user-input { width: calc(100% - 100px); padding: 8px; box-sizing: border-box;}
-        #send-btn { width: 90px; padding: 8px 0;}
-        #input-area { display: flex;}
-    </style>
-</head>
-<body>
-    <h1>Chatbot Profesional (POST API)</h1>
-    <div id="chatbox"></div>
-    <div id="input-area">
-        <input type="text" id="user-input" placeholder="Escribe tu mensaje aquí..." autocomplete="off"/>
-        <button id="send-btn">Enviar</button>
-    </div>
-    <script>
-        const chatbox = document.getElementById('chatbox');
-        const userInput = document.getElementById('user-input');
-        const sendBtn = document.getElementById('send-btn');
-
-        let conversationHistory = [];
-        const DEFAULT_SYSTEM_PROMPT = "Eres un asistente profesional y servicial.";
-        let currentSystemPrompt = DEFAULT_SYSTEM_PROMPT;
-        let botMessageElement = null;
-
-        function appendMessage(sender, text, isStreaming = false) {
-            let msg;
-            if (isStreaming && botMessageElement) {
-                 botMessageElement.textContent += text;
-            } else {
-                msg = document.createElement('p');
-                msg.innerHTML = `<strong>${sender}:</strong> `;
-                const textNode = document.createTextNode(text);
-                msg.appendChild(textNode);
-                chatbox.appendChild(msg);
-                if (sender === 'Bot' && isStreaming) {
-                    botMessageElement = textNode;
-                } else {
-                    botMessageElement = null;
-                }
-            }
-            chatbox.scrollTop = chatbox.scrollHeight;
-        }
-
-        function updateHistory(role, content) {
-             conversationHistory.push({ "role": role, "content": content });
-             const maxHistorySize = 10;
-             if (conversationHistory.length > maxHistorySize * 2) {
-                  conversationHistory = conversationHistory.slice(-(maxHistorySize * 2));
-             }
-        }
-
-        async function sendMessage() {
-            const text = userInput.value;
-            if (!text) {
-                return;
-            }
-            appendMessage('Usuario', text);
-            updateHistory("user", text);
-            userInput.value = '';
-            botMessageElement = null;
-
-            const messagePayload = {
-                input_text: text,
-                history: conversationHistory,
-                system_prompt: currentSystemPrompt,
-                stream: true,
-                temperature: 1.0,
-                top_k: 50,
-                top_p: 1.0,
-                repetition_penalty: 1.0,
-                frequency_penalty: 0.0,
-                presence_penalty: 0.0,
-                num_beams: 1,
-                length_penalty: 1.0,
-                no_repeat_ngram_size: 0,
-                early_stopping: false,
-                do_sample: true,
-                use_mirostat: false,
-                mirostat_tau: 5.0,
-                mirostat_eta: 0.1,
-                max_new_tokens: 512,
-                num_return_sequences: 1,
-                return_token_logprobs: true
-            };
-
-            try {
-                const response = await fetch('/generate', {
-                    method: 'POST',
-                    headers: {
-                        'Content-Type': 'application/json',
-                        // Add API Key header if needed
-                        // 'X-API-Key': 'YOUR_API_KEY_HERE'
-                    },
-                    body: JSON.stringify(messagePayload),
-                });
-
-                if (!response.ok) {
-                    const errorData = await response.json();
-                    throw new Error(`API Error: ${response.status} ${response.statusText} - ${errorData.detail || errorData.error}`);
-                }
-
-                const reader = response.body.getReader();
-                const decoder = new TextDecoder();
-                let buffer = '';
-                let currentBotResponse = "";
-
-                while (true) {
-                    const { value, done } = await reader.read();
-                    if (done) break;
-
-                    buffer += decoder.decode(value, { stream: true });
-
-                    const lines = buffer.split('\n');
-                    buffer = lines.pop();
-
-                    for (const line of lines) {
-                        if (line.trim() === '') continue;
-                        try {
-                            const data = JSON.parse(line);
-                            if (data.type === 'token') {
-                                currentBotResponse += data.text;
-                                appendMessage('Bot', data.text, true);
-                                console.log('Token:', data.token_id, 'Text:', data.text, 'Logprob:', data.logprob);
-                            } else if (data.type === 'done') {
-                                console.log('Generation done', data);
-                                if (data.total_tokens !== undefined) {
-                                     appendMessage('System', `Generated ${data.total_tokens} tokens. Finish reason: ${data.finish_reason}`);
-                                }
-                                if (data.final_text !== undefined) {
-                                     updateHistory("assistant", data.final_text);
-                                } else if (currentBotResponse) {
-                                     updateHistory("assistant", currentBotResponse);
-                                }
-
-                            } else if (data.type === 'error') {
-                                appendMessage('Error', data.message);
-                                currentBotResponse = "";
-                            }
-                        } catch (e) {
-                            console.error('Failed to parse stream chunk:', e, line);
-                            appendMessage('Error', 'Failed to process stream.');
-                            currentBotResponse = "";
-                            reader.cancel();
-                            return;
-                        }
-                    }
-                }
-
-                 if (buffer.trim() !== '') {
-                     try {
-                          const data = JSON.parse(buffer);
-                           if (data.type === 'token') {
-                              currentBotResponse += data.text;
-                              appendMessage('Bot', data.text, true);
-                               console.log('Token:', data.token_id, 'Text:', data.text, 'Logprob:', data.logprob);
-                          } else if (data.type === 'done') {
-                              console.log('Generation done', data);
-                               if (data.total_tokens !== undefined) {
-                                   appendMessage('System', `Generated ${data.total_tokens} tokens. Finish reason: ${data.finish_reason}`);
-                               }
-                               if (data.final_text !== undefined) {
-                                   updateHistory("assistant", data.final_text);
-                               } else if (currentBotResponse) {
-                                   updateHistory("assistant", currentBotResponse);
-                               }
-                          } else if (data.type === 'error') {
-                              appendMessage('Error', data.message);
-                              currentBotResponse = "";
-                          }
-                     } catch (e) {
-                         console.error('Failed to parse remaining buffer:', e, buffer);
-                         appendMessage('Error', 'Failed to process remaining stream data.');
-                         currentBotResponse = "";
-                     }
-                 }
-
-
-                if (currentBotResponse && !botMessageElement) {
-                     updateHistory("assistant", currentBotResponse);
-                }
-                botMessageElement = null;
-                currentBotResponse = "";
-
-
-            } catch (error) {
-                console.error('Send message error:', error);
-                appendMessage('Error', error.message || 'An unknown error occurred.');
-                botMessageElement = null;
-                currentBotResponse = "";
-            }
-        }
-
-        sendBtn.onclick = sendMessage;
-
-        userInput.addEventListener('keypress', function(event) {
-            if (event.key === 'Enter') {
-                event.preventDefault();
-                sendMessage();
-            }
-        });
-
-
-    </script>
-</body>
-</html>
-"""
-
-@app.get("/", response_class=HTMLResponse, summary="Interactive HTML interface")
-async def root():
-    return HTMLResponse(content=html_code)
-
-@app.get("/health", summary="Check API health and model status")
-async def health():
-    if global_model is None or global_tokenizer is None:
-        status_detail = "loading model"
-        if global_tokenizer is None and global_model is not None:
-             status_detail = "tokenizer not loaded"
-        elif global_model is None and global_tokenizer is not None:
-             status_detail = "model not loaded"
-        elif global_model is None and global_tokenizer is None:
-             status_detail = "model and tokenizer not loaded"
-
-        return JSONResponse(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, content={"status": status_detail, "model_loaded": global_model is not None, "tokenizer_loaded": global_tokenizer is not None})
-    return {"status": "ok", "model_loaded": True, "tokenizer_loaded": True}
-
-@app.get("/config", summary="Get API configuration")
-async def get_config():
-    torch_dtype_str_out = str(TORCH_DTYPE).split('.')[-1] if isinstance(TORCH_DTYPE, torch.dtype) else str(TORCH_DTYPE)
-
-    return JSONResponse({
-        "model_name": MODEL_NAME,
-        "system_prompt_default": SYSTEM_PROMPT,
-        "max_context_tokens": MAX_CONTEXT_TOKENS,
-        "max_generation_tokens": MAX_GENERATION_TOKENS,
-        "cuda_available": torch.cuda.is_available(),
-        "model_loaded": global_model is not None,
-        "tokenizer_loaded": global_tokenizer is not None,
-        "max_concurrent_generations": MAX_CONCURRENT_GENERATIONS,
-        "trust_remote_code": TRUST_REMOTE_CODE,
-        "enable_flash_attention_2": ENABLE_FLASH_ATTENTION_2,
-        "torch_dtype": torch_dtype_str_out,
-        "eos_token_id": global_tokens.get("eos_token_id"),
-        "pad_token_id": global_tokens.get("pad_token_id"),
-        "bos_token_id": global_tokenizer.bos_token_id if global_tokenizer else None,
-        "api_key_required": API_KEY is not None
-    })
-
-@app.get("/model/info", summary="Get detailed model information", dependencies=[Depends(get_api_key)])
-async def get_model_info():
-    if global_model is None:
-        raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Model is not loaded.")
-    try:
-        config_dict = global_model.config.to_dict()
-        # Remove potentially large/complex entries if necessary
-        keys_to_remove = ['torch_dtype', '_attn_implementation', 'architectures', 'id2label', 'label2id', 'torch_dtype']
-        for key in keys_to_remove:
-            config_dict.pop(key, None)
-
-        return JSONResponse({
-            "model_name": MODEL_NAME,
-            "is_loaded": True,
-            "device": str(global_model.device),
-            "torch_dtype": str(global_model.dtype),
-            "config": config_dict
-        })
-    except Exception as e:
-        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error getting model info: {e}")
-
-
-@app.post("/tokenize", summary="Tokenize text", dependencies=[Depends(get_api_key)])
-async def tokenize_endpoint(req: TokenizeRequest):
-     if global_tokenizer is None:
-          raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Tokenizer is not loaded.")
-
-     try:
-          tokenizer_kwargs_final = req.tokenizer_kwargs or {}
-
-          return_tensors = req.return_tensors if req.return_tensors is not None else None
-          if return_tensors is None and (req.return_attention_mask or req.return_offsets_mapping or req.return_length):
-               return_tensors = "pt"
-
-          encoded = global_tokenizer(
-               req.text,
-               add_special_tokens=req.add_special_tokens,
-               return_token_type_ids=req.return_token_type_ids,
-               padding=req.padding,
-               truncation=req.truncation,
-               max_length=req.max_length,
-               is_split_into_words=req.is_split_into_words,
-               return_tensors=return_tensors,
-               return_attention_mask=req.return_attention_mask,
-               return_offsets_mapping=req.return_offsets_mapping,
-               return_length=req.return_length,
-               verbose=req.verbose,
-               **tokenizer_kwargs_final
-          )
-          if return_tensors:
-               return JSONResponse({k: v.tolist() if isinstance(v, torch.Tensor) else v for k, v in encoded.items()})
-          else:
-               if isinstance(req.text, str):
-                    response_data = {
-                         "tokens": encoded["input_ids"],
-                         "token_count": len(encoded["input_ids"]),
-                         "input_text": req.text
-                         }
-                    if "token_type_ids" in encoded: response_data["token_type_ids"] = encoded["token_type_ids"]
-                    if "attention_mask" in encoded: response_data["attention_mask"] = encoded["attention_mask"]
-                    if "offset_mapping" in encoded: response_data["offset_mapping"] = encoded["offset_mapping"]
-                    if "length" in encoded: response_data["length"] = encoded["length"]
-                    return JSONResponse(response_data)
-               elif isinstance(req.text, list):
-                    results = []
-                    for i, item_text in enumerate(req.text):
-                         item_result: Dict[str, Any] = {"text": item_text}
-                         if i < len(encoded["input_ids"]):
-                              item_result["tokens"] = encoded["input_ids"][i]
-                              item_result["token_count"] = len(encoded["input_ids"][i])
-                              if "token_type_ids" in encoded and i < len(encoded["token_type_ids"]):
-                                   item_result["token_type_ids"] = encoded["token_type_ids"][i]
-                              if "attention_mask" in encoded and i < len(encoded["attention_mask"]):
-                                   item_result["attention_mask"] = encoded["attention_mask"][i]
-                              if "offset_mapping" in encoded and i < len(encoded["offset_mapping"]):
-                                   item_result["offset_mapping"] = encoded["offset_mapping"][i]
-                              if "length" in encoded and i < len(encoded["length"]):
-                                   item_result["length"] = encoded["length"][i]
-                         else:
-                              item_result["tokens"] = []
-                              item_result["token_count"] = 0
-                         results.append(item_result)
-                    return JSONResponse(results)
-
-
-     except Exception as e:
-          logging.exception("Tokenization error:")
-          raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Tokenization error: {e}")
-
-@app.post("/decode", summary="Decode token IDs to text", dependencies=[Depends(get_api_key)])
-async def decode_endpoint(req: DecodeRequest):
-     if global_tokenizer is None:
-          raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Tokenizer is not loaded.")
-     try:
-          decode_kwargs = req.decode_kwargs or {}
-          text = global_tokenizer.decode(
-              req.token_ids,
-              skip_special_tokens=req.skip_special_tokens,
-              clean_up_tokenization_spaces=req.clean_up_tokenization_spaces,
-              **decode_kwargs
-          )
-          return JSONResponse({"text": text})
-     except Exception as e:
-          logging.exception("Decoding error:")
-          raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Decoding error: {e}")
-
-@app.post("/system_prompt", summary="Update global system prompt", dependencies=[Depends(get_api_key)])
-async def update_system_prompt(req: SystemPromptUpdateRequest):
-     global SYSTEM_PROMPT
-     if req.system_prompt is not None:
-          SYSTEM_PROMPT = req.system_prompt.strip()
-          return JSONResponse({"status": "success", "message": "Global system prompt updated"})
-     else:
-          raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="System prompt cannot be null")
-
-@app.get("/system_prompt", summary="Get global system prompt", dependencies=[Depends(get_api_key)])
-async def get_system_prompt():
-    return JSONResponse({"system_prompt": SYSTEM_PROMPT})
-
-@app.post("/model/reload", summary="Reload the model", dependencies=[Depends(get_api_key)])
-async def reload_model(req: ModelReloadRequest):
-     global global_model, global_tokenizer, global_tokens, MODEL_NAME, TRUST_REMOTE_CODE, ENABLE_FLASH_ATTENTION_2, TORCH_DTYPE, TORCH_DTYPE_STR
-
-     new_model_name = req.model_name if req.model_name else MODEL_NAME
-     new_trust_remote_code = req.trust_remote_code if req.trust_remote_code is not None else TRUST_REMOTE_CODE
-     new_enable_flash_attention_2 = req.enable_flash_attention_2 if req.enable_flash_attention_2 is not None else ENABLE_FLASH_ATTENTION_2
-     new_torch_dtype_str_req = req.torch_dtype if req.torch_dtype else TORCH_DTYPE_STR
-
-     try:
-          new_torch_dtype = getattr(torch, new_torch_dtype_str_req.lower())
-          if new_torch_dtype_str_req.lower() == "bfloat16" and (not torch.cuda.is_available() or (torch.cuda.is_available() and not torch.cuda.get_device_properties(0).has_bfloat16)):
-               logging.warning(f"Requested dtype {new_torch_dtype_str_req} but not supported or CUDA not available. Falling back to float16.")
-               new_torch_dtype = torch.float16
-          elif not isinstance(new_torch_dtype, torch.dtype):
-               raise AttributeError
-          new_torch_dtype_str = str(new_torch_dtype).split('.')[-1]
-     except AttributeError:
-          raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=f"Invalid or unsupported torch_dtype: {new_torch_dtype_str_req}")
-
-
-     device = "cuda" if torch.cuda.is_available() else "cpu"
-
-     async def _reload_model():
-          global global_model, global_tokenizer, global_tokens, MODEL_NAME, TRUST_REMOTE_CODE, ENABLE_FLASH_ATTENTION_2, TORCH_DTYPE, TORCH_DTYPE_STR
-          logging.info(f"Attempting to load model: {new_model_name}")
-          try:
-               logging.info("Unloading current model...")
-               await cleanup(device)
-               if global_model is not None:
-                    del global_model
-                    global_model = None
-               if global_tokenizer is not None:
-                    del global_tokenizer
-                    global_tokenizer = None
-               global_tokens = {}
-               torch.cuda.empty_cache()
-               gc.collect()
-               logging.info("Current model unloaded.")
-
-               logging.info(f"Loading config for model: {new_model_name}")
-               config = AutoConfig.from_pretrained(new_model_name, trust_remote_code=new_trust_remote_code)
-               logging.info("Config loaded.")
-
-               logging.info(f"Loading tokenizer for model: {new_model_name}")
-               tokenizer_kwargs = {"config": config, "trust_remote_code": new_trust_remote_code}
-               if req.tokenizer_kwargs:
-                    tokenizer_kwargs.update(req.tokenizer_kwargs)
-               tokenizer = AutoTokenizer.from_pretrained(new_model_name, **tokenizer_kwargs)
-               logging.info("Tokenizer loaded.")
-
-               logging.info(f"Loading model: {new_model_name} with dtype {new_torch_dtype_str} (CUDA: {torch.cuda.is_available()})")
-               model_kwargs = {"config": config, "torch_dtype": new_torch_dtype, "trust_remote_code": new_trust_remote_code}
-               if new_enable_flash_attention_2 and device == "cuda":
-                    try:
-                         if hasattr(config, '_supports_flash_attn_2') and config._supports_flash_attn_2:
-                             model_kwargs['use_flash_attention_2'] = True
-                             logging.info("Flash Attention 2 enabled for new model.")
-                         elif hasattr(AutoModelForCausalLM, '_supports_flash_attn_2') and AutoModelForCausalLM._supports_flash_attn_2(config):
-                              model_kwargs['use_flash_attention_2'] = True
-                              logging.info("Flash Attention 2 enabled via model check for new model.")
-                         else:
-                             logging.warning("Flash Attention 2 requested but not supported by new model config.")
-                    except Exception as e:
-                         logging.warning(f"Error checking for Flash Attention 2 support for new model: {e}")
-                         pass
-
-               if req.model_kwargs:
-                    model_kwargs.update(req.model_kwargs)
-
-               model = AutoModelForCausalLM.from_pretrained(new_model_name, **model_kwargs)
-               model.to(device)
-               model.eval()
-               logging.info("New model loaded successfully.")
-
-               global_model = model
-               global_tokenizer = tokenizer
-               global_tokens["eos_token_id"] = global_tokenizer.eos_token_id
-               global_tokens["pad_token_id"] = global_tokenizer.pad_token_id
-               if global_tokens["pad_token_id"] is None and global_tokens["eos_token_id"] is not None:
-                    global_tokens["pad_token_id"] = global_tokens["eos_token_id"]
-                    if global_model.config.pad_token_id is None:
-                         global_model.config.pad_token_id = global_tokens["pad_token_id"]
-               elif global_tokens["pad_token_id"] is None and global_tokens["eos_token_id"] is None:
-                    logging.warning("Neither EOS nor PAD token defined for new model.")
-
-
-               MODEL_NAME = new_model_name
-               TRUST_REMOTE_CODE = new_trust_remote_code
-               ENABLE_FLASH_ATTENTION_2 = new_enable_flash_attention_2
-               TORCH_DTYPE = new_torch_dtype
-               TORCH_DTYPE_STR = new_torch_dtype_str
-               if hasattr(global_tokenizer, 'use_fast'):
-                    req.use_fast_tokenizer = global_tokenizer.use_fast
-
-
-               logging.info(f"Model successfully reloaded to: {MODEL_NAME}")
-               logging.info({"status": "success", "message": f"Model {new_model_name} loaded successfully."})
-
-          except Exception as e:
-               logging.exception(f"Failed to load model {new_model_name}:")
-               global_model = None
-               global_tokenizer = None
-               global_tokens = {}
-               logging.error({"status": "error", "message": f"Failed to load model {new_model_name}: {e}. Model is now unloaded."})
-
-
-     asyncio.create_task(_reload_model())
-     return JSONResponse({"status": "info", "message": f"Attempting to load model {new_model_name} in background. Check logs for status."})
-
-@app.post("/model/unload", summary="Unload the model", dependencies=[Depends(get_api_key)])
-async def unload_model():
-     global global_model, global_tokenizer, global_tokens
-     device = "cuda" if torch.cuda.is_available() else "cpu"
-     logging.info("Attempting to unload model.")
-     try:
-          await cleanup(device)
-          if global_model is not None:
-               del global_model
-               global_model = None
-          if global_tokenizer is not None:
-               del global_tokenizer
-               global_tokenizer = None
-          global_tokens = {}
-          torch.cuda.empty_cache()
-          gc.collect()
-          logging.info("Model unloaded successfully.")
-          return JSONResponse({"status": "success", "message": "Model unloaded successfully."})
-     except Exception as e:
-          logging.exception("Failed to unload model:")
-          raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Failed to unload model: {e}")
-
-@app.get("/queue/status", summary="Get current generation queue status", dependencies=[Depends(get_api_key)])
-async def queue_status():
-     try:
-          current_running = MAX_CONCURRENT_GENERATIONS - generation_semaphore._value
-          is_waiting = generation_semaphore.locked()
-
-          return JSONResponse({
-              "max_concurrent": MAX_CONCURRENT_GENERATIONS,
-              "currently_running": current_running,
-              "available_slots": generation_semaphore._value
-          })
-     except Exception as e:
-          logging.exception("Error getting queue status:")
-          raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error getting queue status: {e}")
-
-@app.get("/tokens/special", summary="Get special token IDs", dependencies=[Depends(get_api_key)])
-async def get_special_tokens():
-     if global_tokenizer is None:
-          raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Tokenizer is not loaded.")
-     try:
-          return JSONResponse({
-              "eos_token": global_tokenizer.eos_token,
-              "eos_token_id": global_tokenizer.eos_token_id,
-              "pad_token": global_tokenizer.pad_token,
-              "pad_token_id": global_tokenizer.pad_token_id,
-              "bos_token": global_tokenizer.bos_token,
-              "bos_token_id": global_tokenizer.bos_token_id,
-              "unk_token": global_tokenizer.unk_token,
-              "unk_token_id": global_tokenizer.unk_token_id,
-          })
-     except Exception as e:
-         logging.exception("Error getting special tokens:")
-         raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error getting special tokens: {e}")
-
-@app.post("/generate", summary="Generate text", dependencies=[Depends(get_api_key)])
-async def generate_endpoint(req: GenerateRequest):
-    if global_model is None or global_tokenizer is None:
-        raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Model is not loaded. It may still be loading or failed to load.")
-
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-
-    apply_seed(req.seed)
-
-    try:
-        initial_prompt_text = format_conversation(req.input_text, req.history, req.system_prompt)
-    except Exception as e:
-        logging.exception("Error formatting conversation:")
-        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=f"Error formatting conversation: {e}")
-
-
-    try:
-        tokenizer_encoding_kwargs = req.tokenizer_kwargs or {}
-
-        encoded = global_tokenizer(initial_prompt_text, return_tensors="pt", add_special_tokens=False, **tokenizer_encoding_kwargs).to(device)
-        initial_ids_before_trunc = encoded.input_ids
-        initial_prompt_tokens_count_before_trunc = initial_ids_before_trunc.shape[-1]
-        ids = truncate_encoded_ids(initial_ids_before_trunc, MAX_CONTEXT_TOKENS)
-        current_prompt_tokens_count = ids.shape[-1]
-
-
-    except Exception as e:
-        logging.exception("Tokenizer error during encoding:")
-        await cleanup(device)
-        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Tokenizer encoding error: {e}")
-
-    if req.tokenize_only:
-         await cleanup(device)
-         return JSONResponse({
-             "prompt_tokens_count": initial_prompt_tokens_count_before_trunc,
-             "max_context_tokens": MAX_CONTEXT_TOKENS,
-             "truncated": initial_prompt_tokens_count_before_trunc > MAX_CONTEXT_TOKENS,
-             "input_text_processed": initial_prompt_text,
-             "input_ids_truncated": ids.tolist()[0]
-         })
-
-    total_capacity = MAX_CONTEXT_TOKENS + MAX_GENERATION_TOKENS
-    total_requested_seq_len = current_prompt_tokens_count + req.max_new_tokens
-
-
-    async with generation_semaphore:
-        if req.stream or total_requested_seq_len <= total_capacity:
-
-            gen_cfg = GenerationConfig(
-                temperature=req.temperature,
-                top_k=req.top_k,
-                top_p=req.top_p,
-                repetition_penalty=req.repetition_penalty,
-                frequency_penalty=req.frequency_penalty,
-                presence_penalty=req.presence_penalty,
-                num_beams=req.num_beams if not req.stream else 1,
-                length_penalty=req.length_penalty,
-                no_repeat_ngram_size=req.no_repeat_ngram_size,
-                early_stopping=req.early_stopping,
-                do_sample=req.do_sample,
-                use_mirostat_mode=1 if req.use_mirostat else 0,
-                mirostat_tau=req.mirostat_tau,
-                mirostat_eta=req.mirostat_eta,
-                max_new_tokens=req.max_new_tokens,
-                eos_token_id=req.eos_token_id_override if req.eos_token_id_override is not None else global_tokens.get("eos_token_id"),
-                pad_token_id=req.pad_token_id_override if req.pad_token_id_override is not None else global_tokens.get("pad_token_id"),
-                bos_token_id=req.bos_token_id_override if req.bos_token_id_override is not None else global_tokenizer.bos_token_id,
-                num_return_sequences=req.num_return_sequences if not req.stream else 1,
-                bad_words_ids=req.bad_words_ids,
-                forced_bos_token_id=req.forced_bos_token_id,
-                forced_eos_token_id=req.forced_eos_token_id,
-                renormalize_logits=req.renormalize_logits,
-                suppress_tokens=req.suppress_tokens,
-                begin_suppress_tokens=req.begin_suppress_tokens,
-                end_suppress_tokens=req.end_suppress_tokens,
-                encoder_no_repeat_ngram_size=req.encoder_no_repeat_ngram_size,
-                min_length=req.min_length,
-                max_length=req.max_length,
-                exponential_decay_length_penalty=req.exponential_decay_length_penalty,
-                use_cache=req.use_cache,
-                typical_p=req.typical_p,
-                epsilon_cutoff=req.epsilon_cutoff,
-                eta_cutoff=req.eta_cutoff,
-                temperature_cutoff=req.temperature_cutoff,
-                encoder_repetition_penalty=req.encoder_repetition_penalty,
-                max_time=req.max_time,
-                output_watermark=req.output_watermark,
-                diversity_penalty=req.diversity_penalty,
-                num_beam_groups=req.num_beam_groups if not req.stream else 1,
-                length_normalization_factor=req.length_normalization_factor,
-                min_new_tokens=req.min_new_tokens,
-                do_normalize_logits=req.do_normalize_logits,
-                output_scores=req.output_scores,
-                output_attentions=req.output_attentions,
-                output_hidden_states=req.output_hidden_states,
-            )
-
-            if req.stream:
-                 gen_cfg.use_cache = True
-                 gen_cfg.num_beams = 1
-                 gen_cfg.num_return_sequences = 1
-                 gen_cfg.num_beam_groups = 1
-
-                 try:
-                     json_stream = stream_generation_logic(req, ids, gen_cfg, device)
-                     if req.return_only_text:
-                          return StreamingResponse(plain_text_stream_processor(json_stream, req), media_type="text/plain")
-                     else:
-                         return StreamingResponse(json_stream, media_type="application/json")
-                 except Exception as e:
-                      logging.exception("Error setting up streaming generation:")
-                      await cleanup(device)
-                      raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error setting up streaming generation: {e}")
-
-            else:
-                try:
-                    if total_requested_seq_len <= total_capacity:
-                         response_payload = await non_stream_generation_logic(req, ids, gen_cfg, device)
-                    else:
-                         total_generated_ids = []
-                         original_history = req.history or []
-                         remaining_max_new_tokens = req.max_new_tokens
-                         finish_reason = "max_new_tokens"
-
-                         base_gen_cfg = GenerationConfig(
-                              temperature=req.temperature,
-                              top_k=req.top_k,
-                              top_p=req.top_p,
-                              repetition_penalty=req.repetition_penalty,
-                              frequency_penalty=req.frequency_penalty,
-                              presence_penalty=req.presence_penalty,
-                              num_beams=1,
-                              length_penalty=req.length_penalty,
-                              no_repeat_ngram_size=req.no_repeat_ngram_size,
-                              early_stopping=req.early_stopping,
-                              do_sample=req.do_sample,
-                              use_mirostat_mode=1 if req.use_mirostat else 0,
-                              mirostat_tau=req.mirostat_tau,
-                              mirostat_eta=req.mirostat_eta,
-                              eos_token_id=req.eos_token_id_override if req.eos_token_id_override is not None else global_tokens.get("eos_token_id"),
-                              pad_token_id=req.pad_token_id_override if req.pad_token_id_override is not None else global_tokens.get("pad_token_id"),
-                              bos_token_id=req.bos_token_id_override if req.bos_token_id_override is not None else global_tokenizer.bos_token_id,
-                              num_return_sequences=1,
-                              bad_words_ids=req.bad_words_ids,
-                              forced_bos_token_id=req.forced_bos_token_id,
-                              forced_eos_token_id=req.forced_eos_token_id,
-                              renormalize_logits=req.renormalize_logits,
-                              suppress_tokens=req.suppress_tokens,
-                              begin_suppress_tokens=req.begin_suppress_tokens,
-                              end_suppress_tokens=req.end_suppress_tokens,
-                              encoder_no_repeat_ngram_size=req.encoder_no_repeat_ngram_size,
-                              min_length=0,
-                              max_length=None,
-                              exponential_decay_length_penalty=req.exponential_decay_length_penalty,
-                              use_cache=True,
-                              typical_p=req.typical_p,
-                              epsilon_cutoff=req.epsilon_cutoff,
-                              eta_cutoff=req.eta_cutoff,
-                              temperature_cutoff=req.temperature_cutoff,
-                              encoder_repetition_penalty=req.encoder_repetition_penalty,
-                              max_time=req.max_time,
-                              output_watermark=False,
-                              diversity_penalty=req.diversity_penalty,
-                              num_beam_groups=1,
-                              length_normalization_factor=req.length_normalization_factor,
-                              min_new_tokens=0,
-                              do_normalize_logits=req.do_normalize_logits,
-                              output_scores=False,
-                              output_attentions=False,
-                              output_hidden_states=False,
-                         )
-
-                         while remaining_max_new_tokens > 0:
-                              history_for_this_chunk = copy.deepcopy(original_history)
-                              history_for_this_chunk.append({"role": "user", "content": req.input_text})
-                              if total_generated_ids:
-                                   accumulated_text = global_tokenizer.decode(total_generated_ids, skip_special_tokens=True)
-                                   history_for_this_chunk.append({"role": "assistant", "content": accumulated_text})
-
-                              current_prompt_text = format_conversation(req.input_text if not total_generated_ids else "", history_for_this_chunk, req.system_prompt)
-
-                              encoded_current_prompt = global_tokenizer(current_prompt_text, return_tensors="pt", add_special_tokens=False, **(req.tokenizer_kwargs or {})).to(device)
-                              current_ids = truncate_encoded_ids(encoded_current_prompt.input_ids, MAX_CONTEXT_TOKENS)
-                              current_prompt_tokens_count = current_ids.shape[-1]
-
-                              chunk_max_new_tokens = min(remaining_max_new_tokens, MAX_GENERATION_TOKENS, total_capacity - current_prompt_tokens_count)
-
-                              if chunk_max_new_tokens <= 0:
-                                   finish_reason = "length_limit_reached_no_space"
-                                   break
-
-                              chunk_gen_cfg = copy.deepcopy(base_gen_cfg)
-                              chunk_gen_cfg.max_new_tokens = max(1, chunk_max_new_tokens)
-
-                              stopping_criteria_list = get_stopping_criteria(req, current_ids, global_tokenizer)
-
-                              with torch.no_grad():
-                                   chunk_out = global_model.generate(
-                                      input_ids=current_ids,
-                                      generation_config=chunk_gen_cfg,
-                                      return_dict_in_generate=True,
-                                      output_scores=False,
-                                      output_attentions=False,
-                                      output_hidden_states=False,
-                                      num_return_sequences=1,
-                                      bad_words_ids=req.bad_words_ids,
-                                      suppress_tokens=req.suppress_tokens,
-                                      begin_suppress_tokens=req.begin_suppress_tokens,
-                                      end_suppress_tokens=req.end_suppress_tokens,
-                                      stopping_criteria=stopping_criteria_list if stopping_criteria_list else None,
-                                   )
-
-                              chunk_sequences = chunk_out.sequences
-                              if not chunk_sequences:
-                                   finish_reason = "no_tokens_generated"
-                                   break
-
-                              this_chunk_generated_ids = chunk_sequences[0][current_ids.shape[-1]:].tolist()
-                              total_generated_ids.extend(this_chunk_generated_ids)
-                              remaining_max_new_tokens -= len(this_chunk_generated_ids)
-
-                              eos_token_id = req.eos_token_id_override if req.eos_token_id_override is not None else global_tokens.get("eos_token_id")
-                              if eos_token_id is not None and eos_token_id in this_chunk_generated_ids:
-                                   finish_reason = "eos_token"
-                                   break
-
-                              decoded_this_chunk = global_tokenizer.decode(this_chunk_generated_ids, skip_special_tokens=True)
-                              if req.stop_sequences and any(seq in decoded_this_chunk for seq in req.stop_sequences):
-                                    finish_reason = "stop_sequence"
-                                    break
-
-                              if not this_chunk_generated_ids and remaining_max_new_tokens > 0:
-                                   finish_reason = "stopped_early_no_generation"
-                                   break
-
-                         final_text_raw = global_tokenizer.decode(total_generated_ids, skip_special_tokens=True)
-
-                         if finish_reason == "stop_sequence" and req.stop_sequences:
-                              for stop_seq in req.stop_sequences:
-                                   if stop_seq and stop_seq in final_text_raw:
-                                        final_text_raw = final_text_raw.split(stop_seq, 1)[0]
-                                        break
-
-                         final_text_processed = post_process_text(final_text_raw, req.strip_trailing_whitespace, req.remove_incomplete_sentences)
-
-                         response_payload: Dict[str, Any] = {
-                              "prompt_tokens": initial_prompt_tokens_count_before_trunc,
-                              "generated_sequences": [{
-                                   "text": final_text_processed if req.return_text_from_sequence else None,
-                                   "token_ids": total_generated_ids,
-                                   "generated_tokens_count": len(total_generated_ids),
-                                   "finish_reason": finish_reason,
-                                   "full_sequence_token_ids": (initial_ids_before_trunc.tolist()[0] if initial_ids_before_trunc is not None else []) + total_generated_ids if not req.remove_input_from_output else None,
-                              }],
-                              "total_tokens": initial_prompt_tokens_count_before_trunc + len(total_generated_ids),
-                         }
-
-
-                    if req.return_only_text:
-                         texts = [seq["text"] for seq in response_payload.get("generated_sequences", []) if seq.get("text") is not None]
-                         if req.num_return_sequences == 1 and texts:
-                              return PlainTextResponse(texts[0])
-                         else:
-                              return JSONResponse(texts)
-                    else:
-                         return JSONResponse(response_payload)
-
-                except HTTPException:
-                     raise
-                except Exception as e:
-                    logging.exception("Error during non-streaming generation:")
-                    raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Non-streaming generation error: {e}")
-        else:
-
-             await cleanup(device)
-             raise HTTPException(
-                  status_code=status.HTTP_400_BAD_REQUEST,
-                  detail=f"Requested sequence length ({total_requested_seq_len} tokens = {current_prompt_tokens_count} prompt + {req.max_new_tokens} new) exceeds model capacity ({total_capacity} tokens) and chunking failed or is not applicable/implemented for this configuration."
-              )
-
-
-if __name__ == "__main__":
-    uvicorn.run(
-        app, host="0.0.0.0", port=7860,
-        log_level="critical",
-        access_log=False
-    )