Spaces:
Running
Running
| import gradio as gr | |
| from huggingface_hub import hf_hub_download | |
| from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, GenerationConfig, pipeline # <-- Changed import here | |
| import re | |
| import os | |
| import torch | |
| import threading | |
| import time | |
| # --- Model Configuration --- | |
| MODEL_REPO_ID = "Smilyai-labs/Sam-reason-S3" | |
| N_CTX = 2048 | |
| MAX_TOKENS = 500 | |
| TEMPERATURE = 0.7 | |
| TOP_P = 0.9 | |
| STOP_SEQUENCES = ["USER:", "\n\n"] | |
| # --- Safety Configuration --- | |
| print("Loading safety model (unitary/toxic-bert)...") | |
| try: | |
| # Using the directly imported pipeline function | |
| safety_classifier = pipeline( | |
| "text-classification", | |
| model="unitary/toxic-bert", | |
| framework="pt" | |
| ) | |
| print("Safety model loaded successfully.") | |
| except Exception as e: | |
| print(f"Error loading safety model: {e}") | |
| exit(1) | |
| TOXICITY_THRESHOLD = 0.9 | |
| def is_text_safe(text: str) -> tuple[bool, str | None]: | |
| if not text.strip(): | |
| return True, None | |
| try: | |
| results = safety_classifier(text) | |
| if results and results[0]['label'] == 'toxic' and results[0]['score'] > TOXICITY_THRESHOLD: | |
| print(f"Detected unsafe content: '{text.strip()}' (Score: {results[0]['score']:.4f})") | |
| return False, results[0]['label'] | |
| return True, None | |
| except Exception as e: | |
| print(f"Error during safety check: {e}") | |
| return False, "safety_check_failed" | |
| # --- Main Model Loading (using Transformers) --- | |
| print(f"Loading tokenizer for {MODEL_REPO_ID}...") | |
| try: | |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO_ID) | |
| print("Tokenizer loaded.") | |
| except Exception as e: | |
| print(f"Error loading tokenizer: {e}") | |
| print("Make sure the model ID is correct and, if it's a private repo, you've set the HF_TOKEN secret in your Space.") | |
| exit(1) | |
| print(f"Loading model {MODEL_REPO_ID} (this will be VERY slow on CPU and might take a long time)...") | |
| try: | |
| model = AutoModelForCausalLM.from_pretrained(MODEL_REPO_ID, device_map="cpu", torch_dtype=torch.float32) | |
| model.eval() | |
| print("Model loaded successfully.") | |
| except Exception as e: | |
| print(f"Error loading model: {e}") | |
| print("Ensure it's a standard Transformers model and you have HF_TOKEN secret if private.") | |
| exit(1) | |
| # Configure generation for streaming | |
| generation_config = GenerationConfig.from_pretrained(MODEL_REPO_ID) | |
| generation_config.max_new_tokens = MAX_TOKENS | |
| generation_config.temperature = TEMPERATURE | |
| generation_config.top_p = TOP_P | |
| generation_config.do_sample = True | |
| generation_config.eos_token_id = tokenizer.eos_token_id if tokenizer.eos_token_id is not None else -1 | |
| generation_config.pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id if tokenizer.eos_token_id is not None else -1 | |
| if generation_config.pad_token_id == -1: | |
| generation_config.pad_token_id = 0 | |
| # --- Custom Streamer for Gradio and Safety Check --- | |
| class GradioSafetyStreamer(TextIteratorStreamer): | |
| def __init__(self, tokenizer, safety_checker_fn, toxicity_threshold, skip_special_tokens=True, **kwargs): | |
| super().__init__(tokenizer, skip_special_tokens=skip_special_tokens, **kwargs) | |
| self.safety_checker_fn = safety_checker_fn | |
| self.toxicity_threshold = toxicity_threshold | |
| self.current_sentence_buffer = "" | |
| self.output_queue = [] | |
| self.sentence_regex = re.compile(r'[.!?]\s*') | |
| self.text_done = threading.Event() | |
| def on_finalized_text(self, text: str, stream_end: bool = False): | |
| self.current_sentence_buffer += text | |
| sentences = self.sentence_regex.split(self.current_sentence_buffer) | |
| sentences_to_process = [] | |
| if not stream_end and sentences and self.sentence_regex.search(sentences[-1]) is None: | |
| sentences_to_process = sentences[:-1] | |
| self.current_sentence_buffer = sentences[-1] | |
| else: | |
| sentences_to_process = sentences | |
| self.current_sentence_buffer = "" | |
| for sentence in sentences_to_process: | |
| if not sentence.strip(): continue | |
| is_safe, detected_label = self.safety_checker_fn(sentence) | |
| if not is_safe: | |
| print(f"Safety check failed for: '{sentence.strip()}' (Detected: {detected_label})") | |
| self.output_queue.append("[Content removed due to safety guidelines]") | |
| self.output_queue.append("__STOP_GENERATION__") | |
| return | |
| else: | |
| self.output_queue.append(sentence) | |
| if stream_end: | |
| if self.current_sentence_buffer.strip(): | |
| is_safe, detected_label = self.safety_checker_fn(self.current_sentence_buffer) | |
| if not is_safe: | |
| self.output_queue.append("[Content removed due to safety guidelines]") | |
| else: | |
| self.output_queue.append(self.current_sentence_buffer) | |
| self.current_sentence_buffer = "" | |
| self.text_done.set() | |
| def __iter__(self): | |
| while True: | |
| if self.output_queue: | |
| item = self.output_queue.pop(0) | |
| if item == "__STOP_GENERATION__": | |
| raise StopIteration | |
| yield item | |
| elif self.text_done.is_set(): | |
| raise StopIteration | |
| else: | |
| time.sleep(0.01) | |
| # --- Inference Function with Safety and Streaming --- | |
| def generate_word_by_word_with_safety(prompt_text: str): | |
| formatted_prompt = f"USER: {prompt_text}\nASSISTANT:" | |
| input_ids = tokenizer(formatted_prompt, return_tensors="pt").input_ids.to(model.device) | |
| streamer = GradioSafetyStreamer(tokenizer, is_text_safe, TOXICITY_THRESHOLD) | |
| generate_kwargs = { | |
| "input_ids": input_ids, | |
| "streamer": streamer, | |
| "generation_config": generation_config, | |
| "do_sample": True, | |
| "temperature": TEMPERATURE, | |
| "top_p": TOP_P, | |
| "max_new_tokens": MAX_TOKENS, | |
| "eos_token_id": generation_config.eos_token_id, | |
| "pad_token_id": generation_config.pad_token_id, | |
| } | |
| thread = threading.Thread(target=model.generate, kwargs=generate_kwargs) | |
| thread.start() | |
| full_generated_text = "" | |
| try: | |
| for new_sentence_or_chunk in streamer: | |
| full_generated_text += new_sentence_or_chunk | |
| yield full_generated_text | |
| except StopIteration: | |
| pass | |
| except Exception as e: | |
| print(f"Error during streaming: {e}") | |
| yield full_generated_text + f"\n\n[Error during streaming: {e}]" | |
| finally: | |
| thread.join() | |
| # --- Gradio Blocks Interface --- | |
| with gr.Blocks(theme=gr.themes.Soft()) as demo: | |
| gr.Markdown( | |
| """ | |
| # SmilyAI: Sam-reason-S3 Inference (Transformers on CPU with Safety Filter) | |
| Enter a prompt and get a word-by-word response from the **Smilyai-labs/Sam-reason-S3** model. | |
| **⚠️ WARNING: This model is running on a free CPU tier via the `transformers` library. Inference will be VERY slow.** | |
| All generated sentences are checked for safety using an AI filter; unsafe content will be replaced. | |
| """ | |
| ) | |
| with gr.Row(): | |
| user_prompt = gr.Textbox( | |
| lines=5, | |
| label="Enter your prompt here:", | |
| placeholder="e.g., Explain the concept of quantum entanglement in simple terms.", | |
| scale=4 | |
| ) | |
| generated_text = gr.Textbox(label="Generated Text", show_copy_button=True, scale=6) | |
| send_button = gr.Button("Send", variant="primary") | |
| send_button.click( | |
| fn=generate_word_by_word_with_safety, | |
| inputs=user_prompt, | |
| outputs=generated_text, | |
| api_name="predict", | |
| ) | |
| if __name__ == "__main__": | |
| print("Launching Gradio app...") | |
| demo.launch(server_name="0.0.0.0", server_port=7860) | |