initial commit

gen.py

@@ -1,12 +1,21 @@
 import torch
 import sys
-import sys
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+
+tokenizer = AutoTokenizer.from_pretrained('google/gemma-2-2b-it')
+
+# Configure 4-bit quantization using BitsAndBytesConfig
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.bfloat16,
+    bnb_4bit_quant_type="nf4",
+)
 
-tokenizer = AutoTokenizer.from_pretrained('stabilityai/stablelm-2-zephyr-1_6b')
+# Load the model with the quantization configuration
 model = AutoModelForCausalLM.from_pretrained(
-    'stabilityai/stablelm-2-zephyr-1_6b',
+    'google/gemma-2-2b-it',
     device_map="auto",
+    quantization_config=quantization_config,
 )
 
 
@@ -147,26 +156,43 @@ prompt = (
     "    }\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."
+    "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 empezando con <json>:"
 )
 
 
 def generate(event):
-    # Generar el texto usando el modelo
-    prompt_msg = [{"role":"system","content":prompt},{'role': 'user', 'content': event}]
+    combined_input = f"{prompt} {event}"  # Combine prompt and event
+    prompt_msg = [{'role': 'user', 'content': combined_input}]
+
     inputs = tokenizer.apply_chat_template(
         prompt_msg,
-        add_generation_prompt=False,
+        add_generation_prompt=True,
         return_tensors='pt'
     )
-    
+
     tokens = model.generate(
         inputs.to(model.device),
-        max_new_tokens=20096,
-        temperature=0.7,
+        max_new_tokens=1024,
+        temperature=0.5,
         do_sample=True
     )
-    
 
-    # Imprimir la salida generada
-    return "{".join(tokenizer.decode(tokens[0], skip_special_tokens=True).split("<|user|>")[1].split("{")[1:-1])
+
+    output_text = tokenizer.decode(tokens[0], skip_special_tokens=False)
+    user_prompt_length = len(f"<bos><start_of_turn>user\n{prompt}\n{event}<end_of_turn>\n<start_of_turn>model\n")  # Calculate user prompt length
+
+    json_start_index = output_text.find("<json>")
+    json_end_index = output_text.find("</json>")
+
+    if json_start_index != -1 and json_end_index != -1:
+        json_string = output_text[max(json_start_index + 6, user_prompt_length):json_end_index].strip() # Trim whitespace and remove prompt
+
+        # Validate JSON (you'll need to define a schema for your JSON structure)
+        try:
+            validate(instance=json.loads(json_string), schema=your_json_schema)
+            return json_string
+        except ValidationError as e:
+            return f"Error: Invalid JSON - {e}"
+
+    else:
+        return "Error: <json> or </json> not found in generated output"

psychohistory.py

@@ -11,220 +11,156 @@ def generate_tree(current_x, current_y, depth, max_depth, max_nodes, x_range, G,
     if node_count_per_depth is None:
         node_count_per_depth = {}
 
-    if depth not in node_count_per_depth:
-        node_count_per_depth[depth] = 0
-
     if depth > max_depth:
         return node_count_per_depth
 
+    if depth not in node_count_per_depth:
+        node_count_per_depth[depth] = 0
+
     num_children = random.randint(1, max_nodes)
     x_positions = [current_x + i * x_range / (num_children + 1) for i in range(num_children)]
 
     for x in x_positions:
-        # Add node to the graph
         node_id = len(G.nodes)
         node_count_per_depth[depth] += 1
-        prob = random.uniform(0, 1)  # Assign random probability
-        G.add_node(node_id, pos=(x, prob, depth))  # Use `depth` for z position
+        prob = random.uniform(0, 1)
+        G.add_node(node_id, pos=(x, prob, depth))
         if parent is not None:
             G.add_edge(parent, node_id)
-        # Recursively add child nodes
         generate_tree(x, current_y + 1, depth + 1, max_depth, max_nodes, x_range, G, parent=node_id, node_count_per_depth=node_count_per_depth)
 
     return node_count_per_depth
 
 
-
 def build_graph_from_json(json_data, G):
     """Builds a graph from JSON data."""
+    data = json.loads(json_data)
+
     def add_event(parent_id, event_data, depth):
-        """Recursively adds events and subevents to the graph."""
-        # Add the current event node
         node_id = len(G.nodes)
-        prob = event_data['probability'] / 100.0  # Convert percentage to probability
-        pos = (depth, prob, event_data['event_number'])  # Use event_number for z position
-        label = event_data['name']  # Use event name as label
+        prob = event_data['probability'] / 100.0
+        pos = (depth, prob, event_data['event_number'])
+        label = event_data['name']
         G.add_node(node_id, pos=pos, label=label)
         if parent_id is not None:
             G.add_edge(parent_id, node_id)
 
-        # Add child events
         subevents = event_data.get('subevents', {}).get('event', [])
         if not isinstance(subevents, list):
-            subevents = [subevents]  # Ensure subevents is a list
+            subevents = [subevents]
 
         for subevent in subevents:
             add_event(node_id, subevent, depth + 1)
 
-    data = json.loads(json_data)
-    root_id = len(G.nodes)
     root_event = list(data.get('events', {}).values())[0]
+    root_id = len(G.nodes)
     G.add_node(root_id, pos=(0, root_event['probability'] / 100.0, root_event['event_number']), label=root_event['name'])
-    add_event(None, root_event, 0)  # Start from the root
-
+    add_event(None, root_event, 0)
 
 
 def find_paths(G):
-    """Finds the paths with the highest and lowest average probability, and the longest and shortest durations in graph G."""
-    best_path = None
-    worst_path = None
-    longest_duration_path = None
-    shortest_duration_path = None
-    best_mean_prob = -1
-    worst_mean_prob = float('inf')
-    max_duration = -1
-    min_duration = float('inf')
-
-    for source in G.nodes:
-        for target in G.nodes:
-            if source != target:
-                all_paths = list(nx.all_simple_paths(G, source=source, target=target))
-                for path in all_paths:
-                    # Check if all nodes in the path have the 'pos' attribute
-                    if not all('pos' in G.nodes[node] for node in path):
-                        continue  # Skip paths with nodes missing the 'pos' attribute
-                    
-                    # Calculate the mean probability of the path
-                    probabilities = [G.nodes[node]['pos'][1] for node in path]  # Get node probabilities
-                    mean_prob = np.mean(probabilities)
-                    
-                    # Evaluate path with the highest mean probability
-                    if mean_prob > best_mean_prob:
-                        best_mean_prob = mean_prob
-                        best_path = path
-                    
-                    # Evaluate path with the lowest mean probability
-                    if mean_prob < worst_mean_prob:
-                        worst_mean_prob = mean_prob
-                        worst_path = path
-                    
-                    # Calculate path duration
-                    x_positions = [G.nodes[node]['pos'][0] for node in path]
-                    duration = max(x_positions) - min(x_positions)
-                    
-                    # Evaluate path with the longest duration
-                    if duration > max_duration:
-                        max_duration = duration
-                        longest_duration_path = path
-                    
-                    # Evaluate path with the shortest duration
-                    if duration < min_duration:
-                        min_duration = duration
-                        shortest_duration_path = path
-
-    return best_path, best_mean_prob, worst_path, worst_mean_prob, longest_duration_path, shortest_duration_path
+    """Finds paths with highest/lowest probability and longest/shortest durations."""
+    best_path, worst_path = None, None
+    longest_path, shortest_path = None, None
+    best_mean_prob, worst_mean_prob = -1, float('inf')
+    max_duration, min_duration = -1, float('inf')
+
+    # Use nx.all_pairs_shortest_path for efficiency
+    all_paths_dict = dict(nx.all_pairs_shortest_path(G))
+
+    for source, paths_from_source in all_paths_dict.items():
+        for target, path in paths_from_source.items():
+            if source != target and all('pos' in G.nodes[node] for node in path):
+                probabilities = [G.nodes[node]['pos'][1] for node in path]
+                mean_prob = np.mean(probabilities)
+
+                if mean_prob > best_mean_prob:
+                    best_mean_prob = mean_prob
+                    best_path = path
+                if mean_prob < worst_mean_prob:
+                    worst_mean_prob = mean_prob
+                    worst_path = path
+
+                x_positions = [G.nodes[node]['pos'][0] for node in path]
+                duration = max(x_positions) - min(x_positions)
+
+                if duration > max_duration:
+                    max_duration = duration
+                    longest_path = path
+                if duration < min_duration and duration > 0:  # Avoid paths with 0 duration
+                    min_duration = duration
+                    shortest_path = path
+
+    return best_path, best_mean_prob, worst_path, worst_mean_prob, longest_path, shortest_path
 
 def draw_path_3d(G, path, filename='path_plot_3d.png', highlight_color='blue'):
-    """Draws only the specific path in 3D using networkx and matplotlib and saves the figure to a file."""
-    # Create a subgraph containing only the nodes and edges of the path
+    """Draws a specific path in 3D."""
     H = G.subgraph(path).copy()
-    
     pos = nx.get_node_attributes(G, 'pos')
-    
-    # Get data for 3D visualization
     x_vals, y_vals, z_vals = zip(*[pos[node] for node in path])
-    
+
     fig = plt.figure(figsize=(16, 12))
     ax = fig.add_subplot(111, projection='3d')
 
-    # Assign colors to nodes based on probability
-    node_colors = []
-    for node in path:
-        prob = G.nodes[node]['pos'][1]
-        if prob < 0.33:
-            node_colors.append('red')
-        elif prob < 0.67:
-            node_colors.append('blue')
-        else:
-            node_colors.append('green')
-    
-    # Draw nodes
+    node_colors = ['red' if prob < 0.33 else 'blue' if prob < 0.67 else 'green' for _, prob, _ in [pos[node] for node in path]]
     ax.scatter(x_vals, y_vals, z_vals, c=node_colors, s=700, edgecolors='black', alpha=0.7)
-    
-    # Draw edges
+
     for edge in H.edges():
         x_start, y_start, z_start = pos[edge[0]]
         x_end, y_end, z_end = pos[edge[1]]
         ax.plot([x_start, x_end], [y_start, y_end], [z_start, z_end], color=highlight_color, lw=2)
 
-    # Add labels to nodes
     for node, (x, y, z) in pos.items():
         if node in path:
             ax.text(x, y, z, str(node), fontsize=12, color='black')
 
-    # Set labels and title
     ax.set_xlabel('Time (weeks)')
     ax.set_ylabel('Event Probability')
     ax.set_zlabel('Event Number')
     ax.set_title('3D Event Tree - Path')
 
-    plt.savefig(filename, bbox_inches='tight')  # Save to file with adjusted margins
-    plt.close()  # Close the figure to free resources
+    plt.savefig(filename, bbox_inches='tight')
+    plt.close()
 
 
 def draw_global_tree_3d(G, filename='global_tree.png'):
-    """Draws the entire graph in 3D using networkx and matplotlib and saves the figure to a file."""
+    """Draws the entire graph in 3D."""
     pos = nx.get_node_attributes(G, 'pos')
     labels = nx.get_node_attributes(G, 'label')
-    
-    # Check if the graph is empty
+
     if not pos:
         print("Graph is empty. No nodes to visualize.")
         return
 
-    # Get data for 3D visualization
     x_vals, y_vals, z_vals = zip(*pos.values())
-    
     fig = plt.figure(figsize=(16, 12))
     ax = fig.add_subplot(111, projection='3d')
 
-    # Assign colors to nodes based on probability
-    node_colors = []
-    for node, (x, prob, z) in pos.items():
-        if prob < 0.33:
-            node_colors.append('red')
-        elif prob < 0.67:
-            node_colors.append('blue')
-        else:
-            node_colors.append('green')
-    
-    # Draw nodes
+    node_colors = ['red' if prob < 0.33 else 'blue' if prob < 0.67 else 'green' for _, prob, _ in pos.values()]
     ax.scatter(x_vals, y_vals, z_vals, c=node_colors, s=700, edgecolors='black', alpha=0.7)
-    
-    # Draw edges
+
     for edge in G.edges():
         x_start, y_start, z_start = pos[edge[0]]
         x_end, y_end, z_end = pos[edge[1]]
         ax.plot([x_start, x_end], [y_start, y_end], [z_start, z_end], color='gray', lw=2)
 
-    # Add labels to nodes
     for node, (x, y, z) in pos.items():
         label = labels.get(node, f"{node}")
         ax.text(x, y, z, label, fontsize=12, color='black')
 
-    # Set labels and title
     ax.set_xlabel('Time')
     ax.set_ylabel('Probability')
     ax.set_zlabel('Event Number')
     ax.set_title('3D Event Tree')
 
-    plt.savefig(filename, bbox_inches='tight')  # Save to file with adjusted margins
-    plt.close()  # Close the figure to free resources
+    plt.savefig(filename, bbox_inches='tight')
+    plt.close()
 
 def main(mode, input_file=None):
     G = nx.DiGraph()
 
     if mode == 'random':
-        starting_x = 0
-        starting_y = 0
-        max_depth = 5          # Maximum depth of the tree
-        max_nodes = 3          # Maximum number of child nodes
-        x_range = 10           # Maximum range for x position of nodes
-
-        # Generate the tree and get node count per depth
-        generate_tree(starting_x, starting_y, 0, max_depth, max_nodes, x_range, G)
-
-
+        generate_tree(0, 0, 0, 5, 3, 10, G)
     elif mode == 'json' and input_file:
         with open(input_file, 'r') as file:
             json_data = file.read()
@@ -233,50 +169,33 @@ def main(mode, input_file=None):
         print("Invalid mode or input file not provided.")
         return
 
-    # Save the global visualization
-    draw_global_tree_3d(G, filename='global_tree.png')
-
+    draw_global_tree_3d(G)
 
-    # Find relevant paths
-    best_path, best_mean_prob, worst_path, worst_mean_prob, longest_duration_path, shortest_duration_path = find_paths(G)
+    best_path, best_mean_prob, worst_path, worst_mean_prob, longest_path, shortest_path = find_paths(G)
 
-    # Print results
     if best_path:
-        print(f"\nPath with the highest average probability:")
-        print(" -> ".join(map(str, best_path)))
+        print(f"\nPath with the highest average probability: {' -> '.join(map(str, best_path))}")
         print(f"Average probability: {best_mean_prob:.2f}")
-
     if worst_path:
-        print(f"\nPath with the lowest average probability:")
-        print(" -> ".join(map(str, worst_path)))
+        print(f"\nPath with the lowest average probability: {' -> '.join(map(str, worst_path))}")
         print(f"Average probability: {worst_mean_prob:.2f}")
+    if longest_path:
+        print(f"\nPath with the longest duration: {' -> '.join(map(str, longest_path))}")
+        print(f"Duration: {max(G.nodes[node]['pos'][0] for node in longest_path) - min(G.nodes[node]['pos'][0] for node in longest_path):.2f}")
+    if shortest_path:
+        print(f"\nPath with the shortest duration: {' -> '.join(map(str, shortest_path))}")
+        print(f"Duration: {max(G.nodes[node]['pos'][0] for node in shortest_path) - min(G.nodes[node]['pos'][0] for node in shortest_path):.2f}")
 
-    if longest_duration_path:
-        print(f"\nPath with the longest duration:")
-        print(" -> ".join(map(str, longest_duration_path)))
-        print(f"Duration: {max(G.nodes[node]['pos'][0] for node in longest_duration_path) - min(G.nodes[node]['pos'][0] for node in longest_duration_path):.2f}")
+    draw_global_tree_3d(G)
 
-    if shortest_duration_path:
-        print(f"\nPath with the shortest duration:")
-        print(" -> ".join(map(str, shortest_duration_path)))
-        print(f"Duration: {max(G.nodes[node]['pos'][0] for node in shortest_duration_path) - min(G.nodes[node]['pos'][0] for node in shortest_duration_path):.2f}")
-
-    # Save the global visualization
-    draw_global_tree_3d(G, filename='global_tree.png')
-
-    # Draw and save the 3D figure for each relevant path
     if best_path:
-        draw_path_3d(G, path=best_path, filename='best_path.png', highlight_color='blue')
-
+        draw_path_3d(G, best_path, 'best_path.png', 'blue')
     if worst_path:
-        draw_path_3d(G, path=worst_path, filename='worst_path.png', highlight_color='red')
-
-    if longest_duration_path:
-        draw_path_3d(G, path=longest_duration_path, filename='longest_duration_path.png', highlight_color='green')
-
-    if shortest_duration_path:
-        draw_path_3d(G, path=shortest_duration_path, filename='shortest_duration_path.png', highlight_color='purple')
-
+        draw_path_3d(G, worst_path, 'worst_path.png', 'red')
+    if longest_path:
+        draw_path_3d(G, longest_path, 'longest_duration_path.png', 'green')
+    if shortest_path:
+        draw_path_3d(G, shortest_path, 'shortest_duration_path.png', 'purple')
 
 
 if __name__ == "__main__":
@@ -286,5 +205,3 @@ if __name__ == "__main__":
         mode = sys.argv[1]
         input_file = sys.argv[2] if len(sys.argv) > 2 else None
         main(mode, input_file)
-
-                           

requirements.txt

@@ -1,4 +1,6 @@
-huggingface_hub==0.22.2
+gradio
 torch
+huggingface_hub
 transformers
-accelerate
+accelerate
+bitsandbytes