Spaces:
Sleeping
Sleeping
| import torch | |
| import sys | |
| import sys | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| tokenizer = AutoTokenizer.from_pretrained('stabilityai/stablelm-2-zephyr-1_6b') | |
| model = AutoModelForCausalLM.from_pretrained( | |
| 'stabilityai/stablelm-2-zephyr-1_6b', | |
| device_map="auto", | |
| ) | |
| # Definir el prompt para generar un JSON con eventos anidados | |
| prompt = ( | |
| "Genera un JSON que describa una serie de eventos consecutivos en un formato similar al siguiente:\n\n" | |
| "{\n" | |
| " \"events\": {\n" | |
| " \"event\": {\n" | |
| " \"event_number\": 1,\n" | |
| " \"name\": \"conflict_start\",\n" | |
| " \"description\": \"Tensions escalate between Iran and Israel\",\n" | |
| " \"probability\": 70,\n" | |
| " \"duration_days\": 30,\n" | |
| " \"subevents\": {\n" | |
| " \"event\": {\n" | |
| " \"event_number\": 2,\n" | |
| " \"name\": \"diplomatic_failure\",\n" | |
| " \"description\": \"Diplomatic negotiations fail\",\n" | |
| " \"probability\": 60,\n" | |
| " \"duration_days\": 15,\n" | |
| " \"subevents\": {\n" | |
| " \"event\": {\n" | |
| " \"event_number\": 3,\n" | |
| " \"name\": \"military_clash\",\n" | |
| " \"description\": \"Initial military clash at the border\",\n" | |
| " \"probability\": 50,\n" | |
| " \"duration_days\": 10,\n" | |
| " \"subevents\": {\n" | |
| " \"event\": [\n" | |
| " {\n" | |
| " \"event_number\": 4,\n" | |
| " \"name\": \"escalation\",\n" | |
| " \"description\": \"Conflict escalates into full-scale war\",\n" | |
| " \"probability\": 40,\n" | |
| " \"duration_days\": 180,\n" | |
| " \"subevents\": {\n" | |
| " \"event\": [\n" | |
| " {\n" | |
| " \"event_number\": 5,\n" | |
| " \"name\": \"regional_involvement\",\n" | |
| " \"description\": \"Other Middle Eastern countries get involved\",\n" | |
| " \"probability\": 30,\n" | |
| " \"duration_days\": 365,\n" | |
| " \"subevents\": {\n" | |
| " \"event\": [\n" | |
| " {\n" | |
| " \"event_number\": 6,\n" | |
| " \"name\": \"ceasefire\",\n" | |
| " \"description\": \"International powers broker a ceasefire\",\n" | |
| " \"probability\": 20,\n" | |
| " \"duration_days\": 30\n" | |
| " },\n" | |
| " {\n" | |
| " \"event_number\": 7,\n" | |
| " \"name\": \"prolonged_conflict\",\n" | |
| " \"description\": \"Conflict continues for over a year\",\n" | |
| " \"probability\": 50,\n" | |
| " \"duration_days\": 365\n" | |
| " }\n" | |
| " ]\n" | |
| " }\n" | |
| " },\n" | |
| " {\n" | |
| " \"event_number\": 8,\n" | |
| " \"name\": \"international_intervention\",\n" | |
| " \"description\": \"UN or other international organizations intervene\",\n" | |
| " \"probability\": 25,\n" | |
| " \"duration_days\": 60\n" | |
| " }\n" | |
| " ]\n" | |
| " }\n" | |
| " },\n" | |
| " {\n" | |
| " \"event_number\": 9,\n" | |
| " \"name\": \"containment\",\n" | |
| " \"description\": \"Conflict is contained and doesn't escalate\",\n" | |
| " \"probability\": 30,\n" | |
| " \"duration_days\": 90\n" | |
| " }\n" | |
| " ]\n" | |
| " }\n" | |
| " },\n" | |
| " \"event\": {\n" | |
| " \"event_number\": 10,\n" | |
| " \"name\": \"sanctions\",\n" | |
| " \"description\": \"Increased sanctions on Iran\",\n" | |
| " \"probability\": 70,\n" | |
| " \"duration_days\": 180,\n" | |
| " \"subevents\": {\n" | |
| " \"event\": [\n" | |
| " {\n" | |
| " \"event_number\": 11,\n" | |
| " \"name\": \"iran_retaliates\",\n" | |
| " \"description\": \"Iran retaliates with cyberattacks\",\n" | |
| " \"probability\": 40,\n" | |
| " \"duration_days\": 60\n" | |
| " },\n" | |
| " {\n" | |
| " \"event_number\": 12,\n" | |
| " \"name\": \"israel_response\",\n" | |
| " \"description\": \"Israel responds with targeted airstrikes\",\n" | |
| " \"probability\": 50,\n" | |
| " \"duration_days\": 60\n" | |
| " }\n" | |
| " ]\n" | |
| " }\n" | |
| " }\n" | |
| " }\n" | |
| " },\n" | |
| " \"event\": {\n" | |
| " \"event_number\": 13,\n" | |
| " \"name\": \"diplomatic_success\",\n" | |
| " \"description\": \"Successful diplomatic negotiations\",\n" | |
| " \"probability\": 40,\n" | |
| " \"duration_days\": 30,\n" | |
| " \"subevents\": {\n" | |
| " \"event\": [\n" | |
| " {\n" | |
| " \"event_number\": 14,\n" | |
| " \"name\": \"peace_agreement\",\n" | |
| " \"description\": \"Iran and Israel sign a peace agreement\",\n" | |
| " \"probability\": 20,\n" | |
| " \"duration_days\": 60\n" | |
| " },\n" | |
| " {\n" | |
| " \"event_number\": 15,\n" | |
| " \"name\": \"temporary_truce\",\n" | |
| " \"description\": \"A temporary truce is established\",\n" | |
| " \"probability\": 30,\n" | |
| " \"duration_days\": 30\n" | |
| " }\n" | |
| " ]\n" | |
| " }\n" | |
| " }\n" | |
| " }\n" | |
| " }\n" | |
| " }\n" | |
| "}\n\n" | |
| "Ahora, genera un JSON similar con eventos anidados, pero cambia los detalles y números para hacer que sea con el input que viene a continuacion, respondiendo solo el JSON. No muestres el mensaje del sistema del rol system. Contesta solo JSON, el JSON de respuesta.No muestres este mensaje ni el prompt introducido por el usuario. Asegurate en la respuesta que el JSON esta completo y tiene el formato correcto." | |
| ) | |
| def generate(event): | |
| # Generar el texto usando el modelo | |
| prompt_msg = [{"role":"system","content":prompt},{'role': 'user', 'content': event}] | |
| inputs = tokenizer.apply_chat_template( | |
| prompt_msg, | |
| add_generation_prompt=False, | |
| return_tensors='pt' | |
| ) | |
| tokens = model.generate( | |
| inputs.to(model.device), | |
| max_new_tokens=20096, | |
| temperature=0.7, | |
| do_sample=True | |
| ) | |
| # Imprimir la salida generada | |
| return "{".join(tokenizer.decode(tokens[0], skip_special_tokens=True).split("<|user|>")[1].split("{")[1:-1]) | |