Polarium/NextTokenPrediction
1
1import gradio as gr2import torch3from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForSeq2SeqLM4import json5from typing import Dict, List, Tuple6import numpy as np7 8# Global variables for models9device = "cuda" if torch.cuda.is_available() else "cpu"10print(f"Using device: {device}")11 12# Model names13TEXT_GEN_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"14SUMMARIZATION_MODEL = "facebook/bart-large-cnn"15 16# Load models and tokenizers17print("Loading models...")18gen_tokenizer = AutoTokenizer.from_pretrained(TEXT_GEN_MODEL)19gen_model = AutoModelForCausalLM.from_pretrained(TEXT_GEN_MODEL).to(device)20 21sum_tokenizer = AutoTokenizer.from_pretrained(SUMMARIZATION_MODEL)22sum_model = AutoModelForSeq2SeqLM.from_pretrained(SUMMARIZATION_MODEL).to(device)23print("Models loaded successfully!")24 25 26def count_words(text: str) -> int:27 """Count words in text"""28 return len(text.split())29 30 31def generate_text_with_alternatives(32 input_text: str,33 max_tokens: int = 10034) -> Tuple[str, List[Dict]]:35 """36 Generate text and capture top-5 alternative tokens for each generated token.37 Returns: (generated_text, token_alternatives)38 """39 # Prepare input40 messages = [{"role": "user", "content": input_text}]41 text = gen_tokenizer.apply_chat_template(42 messages,43 tokenize=False,44 add_generation_prompt=True45 )46 inputs = gen_tokenizer(text, return_tensors="pt").to(device)47 48 # Generate with output_scores to get token probabilities49 with torch.no_grad():50 outputs = gen_model.generate(51 **inputs,52 max_new_tokens=max_tokens,53 output_scores=True,54 return_dict_in_generate=True,55 do_sample=False, # Greedy decoding56 pad_token_id=gen_tokenizer.eos_token_id57 )58 59 # Get generated tokens (excluding input)60 generated_ids = outputs.sequences[0][inputs.input_ids.shape[1]:]61 generated_text = gen_tokenizer.decode(generated_ids, skip_special_tokens=True)62 63 # Extract token alternatives from scores64 token_alternatives = []65 if hasattr(outputs, 'scores') and outputs.scores:66 for score_tensor in outputs.scores:67 # Get probabilities68 probs = torch.nn.functional.softmax(score_tensor[0], dim=-1)69 70 # Get top 5 tokens71 top_probs, top_indices = torch.topk(probs, k=5)72 73 alternatives = []74 for prob, idx in zip(top_probs, top_indices):75 token = gen_tokenizer.decode([idx.item()])76 alternatives.append({77 "token": token,78 "probability": f"{prob.item() * 100:.2f}%"79 })80 81 token_alternatives.append(alternatives)82 83 return generated_text, token_alternatives84 85 86def summarize_text_with_alternatives(87 input_text: str,88 max_tokens: int = 10089) -> Tuple[str, List[Dict]]:90 """91 Summarize text and capture top-5 alternative tokens for each generated token.92 Returns: (summary_text, token_alternatives)93 """94 inputs = sum_tokenizer(input_text, return_tensors="pt", max_length=1024, truncation=True).to(device)95 96 # Generate with output_scores97 with torch.no_grad():98 outputs = sum_model.generate(99 **inputs,100 max_length=max_tokens,101 output_scores=True,102 return_dict_in_generate=True,103 do_sample=False, # Greedy decoding104 )105 106 # Decode summary107 summary_text = sum_tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)108 109 # Extract token alternatives110 token_alternatives = []111 if hasattr(outputs, 'scores') and outputs.scores:112 for score_tensor in outputs.scores:113 probs = torch.nn.functional.softmax(score_tensor[0], dim=-1)114 top_probs, top_indices = torch.topk(probs, k=5)115 116 alternatives = []117 for prob, idx in zip(top_probs, top_indices):118 token = sum_tokenizer.decode([idx.item()])119 alternatives.append({120 "token": token,121 "probability": f"{prob.item() * 100:.2f}%"122 })123 124 token_alternatives.append(alternatives)125 126 return summary_text, token_alternatives127 128 129def create_html_with_tooltips(text: str, token_alternatives: List[Dict]) -> str:130 """131 Create HTML with hoverable words that show token alternatives.132 """133 if not token_alternatives:134 return f"<div style='padding: 20px; font-size: 16px;'>{text}</div>"135 136 # Split text into tokens/words for display137 words = text.split()138 139 html_parts = []140 html_parts.append("""141 <style>142 .word-container {143 display: inline-block;144 position: relative;145 margin: 2px;146 padding: 2px 4px;147 cursor: pointer;148 border-radius: 3px;149 transition: background-color 0.2s;150 }151 .word-container:hover {152 background-color: #e3f2fd;153 }154 .tooltip {155 visibility: hidden;156 position: absolute;157 z-index: 1000;158 background-color: #263238;159 color: white;160 padding: 12px;161 border-radius: 6px;162 font-size: 13px;163 min-width: 250px;164 bottom: 125%;165 left: 50%;166 transform: translateX(-50%);167 box-shadow: 0 4px 6px rgba(0,0,0,0.3);168 opacity: 0;169 transition: opacity 0.3s;170 }171 .tooltip::after {172 content: "";173 position: absolute;174 top: 100%;175 left: 50%;176 margin-left: -5px;177 border-width: 5px;178 border-style: solid;179 border-color: #263238 transparent transparent transparent;180 }181 .word-container:hover .tooltip {182 visibility: visible;183 opacity: 1;184 }185 .alternative-item {186 padding: 4px 0;187 border-bottom: 1px solid #37474f;188 }189 .alternative-item:last-child {190 border-bottom: none;191 }192 .token-text {193 font-weight: bold;194 color: #81d4fa;195 }196 .probability {197 float: right;198 color: #a5d6a7;199 }200 .result-container {201 padding: 20px;202 font-size: 16px;203 line-height: 1.8;204 font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif;205 }206 </style>207 <div class='result-container'>208 """)209 210 # Map words to token alternatives (approximate mapping)211 alt_index = 0212 for word in words:213 if alt_index < len(token_alternatives):214 alternatives = token_alternatives[alt_index]215 216 # Create tooltip content217 tooltip_html = "<div class='tooltip'>"218 tooltip_html += "<div style='margin-bottom: 8px; font-weight: bold; border-bottom: 2px solid #37474f; padding-bottom: 4px;'>Top 5 Alternatives:</div>"219 for i, alt in enumerate(alternatives, 1):220 tooltip_html += f"<div class='alternative-item'>"221 tooltip_html += f"<span>{i}. <span class='token-text'>{alt['token']}</span></span>"222 tooltip_html += f"<span class='probability'>{alt['probability']}</span>"223 tooltip_html += f"</div>"224 tooltip_html += "</div>"225 226 html_parts.append(f"<span class='word-container'>{word}{tooltip_html}</span>")227 alt_index += 1228 else:229 html_parts.append(f"<span class='word-container'>{word}</span>")230 231 html_parts.append("</div>")232 return "".join(html_parts)233 234 235def process_text(input_text: str, mode: str, max_tokens: int) -> Tuple[str, str]:236 """237 Main processing function that handles both text generation and summarization.238 Returns: (result_html, status_message)239 """240 if not input_text or not input_text.strip():241 return "<div style='padding: 20px; color: red;'>Please enter some text to process.</div>", "❌ No input provided"242 243 # Check word count244 word_count = count_words(input_text)245 if word_count > 500:246 return f"<div style='padding: 20px; color: red;'>Input exceeds maximum limit of 500 words. Current: {word_count} words.</div>", f"❌ Input too long ({word_count} words)"247 248 try:249 if mode == "Text Generation":250 status = f"🔄 Generating text (max {max_tokens} tokens)..."251 generated_text, alternatives = generate_text_with_alternatives(input_text, max_tokens)252 result_html = create_html_with_tooltips(generated_text, alternatives)253 return result_html, f"✅ Generated {len(alternatives)} tokens"254 else: # Text Summarization255 status = f"🔄 Summarizing text (max {max_tokens} tokens)..."256 summary_text, alternatives = summarize_text_with_alternatives(input_text, max_tokens)257 result_html = create_html_with_tooltips(summary_text, alternatives)258 return result_html, f"✅ Generated {len(alternatives)} tokens"259 except Exception as e:260 error_msg = f"<div style='padding: 20px; color: red;'>Error: {str(e)}</div>"261 return error_msg, f"❌ Error: {str(e)}"262 263 264# Create Gradio interface265with gr.Blocks(title="AI Text Assistant", theme=gr.themes.Soft()) as demo:266 gr.Markdown("""267 # 🤖 AI Text Assistant268 Generate text or summarize articles using state-of-the-art AI models.269 **Hover over any word** in the result to see the top 5 alternative tokens the AI considered!270 """)271 272 with gr.Row():273 with gr.Column(scale=2):274 mode = gr.Radio(275 choices=["Text Generation", "Text Summarization"],276 value="Text Generation",277 label="Mode",278 info="Choose between generating new text or summarizing existing text"279 )280 281 input_text = gr.Textbox(282 label="Input Text",283 placeholder="Enter your text here... (max 500 words)",284 lines=6,285 max_lines=10286 )287 288 with gr.Row():289 max_tokens = gr.Slider(290 minimum=10,291 maximum=500,292 value=100,293 step=10,294 label="Max Tokens",295 info="Maximum number of tokens to generate"296 )297 298 process_btn = gr.Button("🚀 Process", variant="primary", size="lg")299 status = gr.Textbox(label="Status", interactive=False)300 301 with gr.Row():302 output_html = gr.HTML(label="Result")303 304 gr.Markdown("""305 ### 💡 Tips:306 - **Text Generation**: Provide a prompt and the AI will continue writing307 - **Text Summarization**: Paste an article or long text to get a concise summary308 - **Hover** over any word in the output to see what other words the AI considered309 - Models used: Qwen/Qwen2.5-0.5B-Instruct (generation) & facebook/bart-large-cnn (summarization)310 """)311 312 # Connect the button to the processing function313 process_btn.click(314 fn=process_text,315 inputs=[input_text, mode, max_tokens],316 outputs=[output_html, status]317 )318 319if __name__ == "__main__":320 demo.launch()321