import torch, torch.nn.functional as F from tqdm import trange import time from tokenxxx import * from main import * from duckduckgo_search import DDGS try: END_OF_TEXT_TOKEN except NameError: END_OF_TEXT_TOKEN = "" try: SYSTEM_PROMPT except NameError: SYSTEM_PROMPT = "Sistema: Proporcione respuestas ultra rápidas, coherentes, similares, precisas y con sentido, con razonamiento lógico y profundo." try: MAX_XDD except NameError: MAX_XDD = 5 try: codegen_model except NameError: codegen_model = None try: codegen_tokenizer except NameError: codegen_tokenizer = None try: summarization_model except NameError: summarization_model = None try: summarization_tokenizer except NameError: summarization_tokenizer = None try: model_gpt2 except NameError: model_gpt2 = None try: enc except NameError: enc = None try: device except NameError: device = "cpu" if torch.device(device).type == "cuda": torch.backends.cudnn.benchmark = True MAX_GENERATION_LENGTH = 512 def top_k_top_p_filtering(logits, top_k=0, top_p=0.0, filter_value=-float('Inf')): top_k = min(top_k, logits.size(-1)); if top_k > 0: indices_to_remove = logits < torch.topk(logits, top_k)[0][..., [-1]]; logits[indices_to_remove] = filter_value if top_p > 0.0: sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1); cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) sorted_indices_to_remove = cumulative_probs > top_p; sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone(); sorted_indices_to_remove[..., 0] = 0; indices_to_remove = sorted_indices[sorted_indices_to_remove]; logits[indices_to_remove] = filter_value; return logits def _generate_sequence(model_call, context_tensor, generated, decode_fn, end_token_condition, temperature, top_k, top_p, repetition_penalty, max_length): past_key_values = None; last_token = None; repetition_count = 0 for _ in range(max_length): try: outputs = model_call(context_tensor, past_key_values) except Exception as e: yield ""; yield ""; return next_token_logits = outputs[0][:, -1, :] / temperature; past_key_values = outputs[1] for token_index in set(generated): next_token_logits[0, token_index] /= repetition_penalty filtered_logits = top_k_top_p_filtering(next_token_logits, top_k=top_k, top_p=top_p) if temperature == 0: next_token = torch.argmax(filtered_logits, dim=-1).unsqueeze(0) else: next_token = torch.multinomial(F.softmax(filtered_logits, dim=-1), num_samples=1) token_id = next_token.tolist()[0][0] if token_id == last_token: repetition_count += 1 else: repetition_count = 0; last_token = token_id if repetition_count >= 10: yield ""; return generated.append(token_id); token_decoded = decode_fn(token_id); yield token_decoded if end_token_condition(token_id): yield ""; return def sample_sequence(prompt, model, enc, max_length=MAX_GENERATION_LENGTH, temperature=1, top_k=0, top_p=0.0, repetition_penalty=1.0, device="cpu"): context_tokens = enc.encode(prompt); context_tensor = torch.tensor([context_tokens], dtype=torch.long, device=device) return _generate_sequence(lambda ct, past: model(ct, past_key_values=past), context_tensor, list(context_tokens), lambda token: enc.decode([token]), lambda token: token == enc.encoder[END_OF_TEXT_TOKEN], temperature, top_k, top_p, repetition_penalty, max_length) def sample_sequence_codegen(prompt, model, tokenizer, max_length=MAX_GENERATION_LENGTH, temperature=1, top_k=0, top_p=0.0, repetition_penalty=1.0, device="cpu"): context_tokens = tokenizer.encode(prompt); context_tensor = torch.tensor([context_tokens], dtype=torch.long, device=device) return _generate_sequence(lambda ct, past: model(input_ids=ct, past_key_values=past, labels=None), context_tensor, list(context_tokens), lambda token: tokenizer.decode([token]), lambda token: token == 50256, temperature, top_k, top_p, repetition_penalty, max_length) def summarize_text(text): if summarization_model and summarization_tokenizer: input_ids = summarization_tokenizer.encode(text, return_tensors="pt", truncation=True, max_length=1024).to(device); summary_ids = summarization_model.generate(input_ids, max_length=150, min_length=40, length_penalty=2.0, num_beams=4, early_stopping=True) return summarization_tokenizer.decode(summary_ids[0], skip_special_tokens=True) return text[:300] + "..." if len(text) > 300 else text def perform_reasoning_stream(text_input, temperature, top_k, top_p, repetition_penalty, prev_context=""): initial_prompt = SYSTEM_PROMPT + "\n\nUser: " + text_input + "\nAssistant:"; reasoning_prompt = prev_context if prev_context else initial_prompt; ddgs = DDGS() search_results = [r for r in ddgs.text(text_input, max_results=MAX_XDD)]; if search_results: reasoning_prompt += "\nWeb Search Results:\n"; for result in search_results: reasoning_prompt += "- " + result['body'] + "\n" reasoning_prompt += "\n" if "code" in text_input.lower() or "program" in text_input.lower(): model_type = "code" elif "summarize" in text_input.lower() or "summary" in text_input.lower(): model_type = "summarize" elif model_gpt2 and enc: model_type = "gpt2" else: yield ""; yield ""; return if model_type == "summarize": if summarization_model: summary = summarize_text(text_input); yield "SUMMARY_TEXT:" + summary; yield ""; return accumulated_text = ""; current_context = reasoning_prompt; overlap = 256 while True: if model_type == "code": generator = sample_sequence_codegen(current_context, codegen_model, codegen_tokenizer, MAX_GENERATION_LENGTH, temperature, top_k, top_p, repetition_penalty, device) elif model_type == "gpt2": generator = sample_sequence(current_context, model_gpt2, enc, MAX_GENERATION_LENGTH, temperature, top_k, top_p, repetition_penalty, device) chunk_text = ""; finished = False for token in generator: if token == "": finished = True; break chunk_text += token if accumulated_text: overlap_text = accumulated_text[-overlap:]; if chunk_text.startswith(overlap_text): chunk_text = chunk_text[len(overlap_text):] accumulated_text += chunk_text; yield chunk_text if finished: yield ""; break current_context = accumulated_text[-overlap:] if len(accumulated_text) > overlap else accumulated_text