Upload demo_linux_fc.py with huggingface_hub

demo_linux_fc.py

@@ -0,0 +1,594 @@
+#!/usr/bin/env python3
+"""
+HybriKo-117M Linux Function Calling Demo
+Usage: python scripts/demo_linux_fc.py
+
+Example:
+  [사용자] 현재 폴더의 파일 목록을 보여줘
+  [HybriKo] Thought: 디렉토리 내용을 확인합니다.
+           Action: ls_command
+           Action Input: {"path": ".", "options": "-l"}
+"""
+
+import torch
+import sentencepiece as spm
+import sys
+import json
+import re
+import argparse
+
+sys.path.insert(0, ".")
+from hybridko.model import HybriKoModel, HybriKoConfig
+
+
+# Exact system prompt used during training (with prettified JSON)
+SYSTEM_PROMPT = """You are a Linux command assistant. You can use many tools (functions) to help users with their Linux tasks.
+At each step, you need to give your thought to analyze the status now and what to do next, with a function call to actually execute your step. Your output should follow this format:
+Thought:
+Action
+Action Input:
+
+After the call, you will get the call result, and you are now in a new state.
+Then you will analyze your status now, then decide what to do next...
+After many (Thought-call) pairs, you finally perform the task, then you can give your final answer.
+
+Remember:
+1. The state change is irreversible, you can't go back to one of the former state.
+2. All the thought is short, at most in 5 sentences.
+3. ALWAYS call "Finish" function at the end of the task.
+4. If you cannot handle the task with the available tools, say you don't know and call Finish with give_answer.
+
+You have access of the following tools:
+[
+  {
+    "name": "ls_command",
+    "description": "List directory contents.",
+    "parameters": {
+      "type": "object",
+      "properties": {
+        "path": {
+          "type": "string"
+        },
+        "options": {
+          "type": "string"
+        }
+      },
+      "required": [
+        "path"
+      ]
+    }
+  },
+  {
+    "name": "cd_command",
+    "description": "Change the current working directory.",
+    "parameters": {
+      "type": "object",
+      "properties": {
+        "path": {
+          "type": "string"
+        }
+      },
+      "required": [
+        "path"
+      ]
+    }
+  },
+  {
+    "name": "mkdir_command",
+    "description": "Create a new directory.",
+    "parameters": {
+      "type": "object",
+      "properties": {
+        "path": {
+          "type": "string"
+        }
+      },
+      "required": [
+        "path"
+      ]
+    }
+  },
+  {
+    "name": "rm_command",
+    "description": "Remove files or directories.",
+    "parameters": {
+      "type": "object",
+      "properties": {
+        "path": {
+          "type": "string"
+        },
+        "recursive": {
+          "type": "boolean"
+        }
+      },
+      "required": [
+        "path"
+      ]
+    }
+  },
+  {
+    "name": "cp_command",
+    "description": "Copy files or directories.",
+    "parameters": {
+      "type": "object",
+      "properties": {
+        "source": {
+          "type": "string"
+        },
+        "destination": {
+          "type": "string"
+        }
+      },
+      "required": [
+        "source",
+        "destination"
+      ]
+    }
+  },
+  {
+    "name": "mv_command",
+    "description": "Move or rename files.",
+    "parameters": {
+      "type": "object",
+      "properties": {
+        "source": {
+          "type": "string"
+        },
+        "destination": {
+          "type": "string"
+        }
+      },
+      "required": [
+        "source",
+        "destination"
+      ]
+    }
+  },
+  {
+    "name": "find_command",
+    "description": "Find files by name pattern.",
+    "parameters": {
+      "type": "object",
+      "properties": {
+        "path": {
+          "type": "string"
+        },
+        "name": {
+          "type": "string"
+        }
+      },
+      "required": [
+        "path",
+        "name"
+      ]
+    }
+  },
+  {
+    "name": "cat_command",
+    "description": "Display file contents.",
+    "parameters": {
+      "type": "object",
+      "properties": {
+        "file": {
+          "type": "string"
+        }
+      },
+      "required": [
+        "file"
+      ]
+    }
+  },
+  {
+    "name": "grep_command",
+    "description": "Search for patterns in files.",
+    "parameters": {
+      "type": "object",
+      "properties": {
+        "pattern": {
+          "type": "string"
+        },
+        "file": {
+          "type": "string"
+        }
+      },
+      "required": [
+        "pattern",
+        "file"
+      ]
+    }
+  },
+  {
+    "name": "head_command",
+    "description": "Display first lines of a file.",
+    "parameters": {
+      "type": "object",
+      "properties": {
+        "file": {
+          "type": "string"
+        },
+        "lines": {
+          "type": "integer"
+        }
+      },
+      "required": [
+        "file"
+      ]
+    }
+  },
+  {
+    "name": "tail_command",
+    "description": "Display last lines of a file.",
+    "parameters": {
+      "type": "object",
+      "properties": {
+        "file": {
+          "type": "string"
+        },
+        "lines": {
+          "type": "integer"
+        }
+      },
+      "required": [
+        "file"
+      ]
+    }
+  },
+  {
+    "name": "wc_command",
+    "description": "Count lines, words, and bytes.",
+    "parameters": {
+      "type": "object",
+      "properties": {
+        "file": {
+          "type": "string"
+        }
+      },
+      "required": [
+        "file"
+      ]
+    }
+  },
+  {
+    "name": "ps_command",
+    "description": "Display running processes.",
+    "parameters": {
+      "type": "object",
+      "properties": {
+        "options": {
+          "type": "string"
+        }
+      },
+      "required": []
+    }
+  },
+  {
+    "name": "df_command",
+    "description": "Display disk space usage.",
+    "parameters": {
+      "type": "object",
+      "properties": {
+        "options": {
+          "type": "string"
+        }
+      },
+      "required": []
+    }
+  },
+  {
+    "name": "du_command",
+    "description": "Display directory space usage.",
+    "parameters": {
+      "type": "object",
+      "properties": {
+        "path": {
+          "type": "string"
+        },
+        "options": {
+          "type": "string"
+        }
+      },
+      "required": [
+        "path"
+      ]
+    }
+  },
+  {
+    "name": "top_command",
+    "description": "Display system processes in real-time.",
+    "parameters": {
+      "type": "object",
+      "properties": {},
+      "required": []
+    }
+  },
+  {
+    "name": "ping_command",
+    "description": "Test network connectivity.",
+    "parameters": {
+      "type": "object",
+      "properties": {
+        "host": {
+          "type": "string"
+        },
+        "count": {
+          "type": "integer"
+        }
+      },
+      "required": [
+        "host"
+      ]
+    }
+  },
+  {
+    "name": "curl_command",
+    "description": "Transfer data from URL.",
+    "parameters": {
+      "type": "object",
+      "properties": {
+        "url": {
+          "type": "string"
+        },
+        "options": {
+          "type": "string"
+        }
+      },
+      "required": [
+        "url"
+      ]
+    }
+  },
+  {
+    "name": "chmod_command",
+    "description": "Change file permissions.",
+    "parameters": {
+      "type": "object",
+      "properties": {
+        "mode": {
+          "type": "string"
+        },
+        "file": {
+          "type": "string"
+        }
+      },
+      "required": [
+        "mode",
+        "file"
+      ]
+    }
+  },
+  {
+    "name": "tar_command",
+    "description": "Archive or extract files.",
+    "parameters": {
+      "type": "object",
+      "properties": {
+        "options": {
+          "type": "string"
+        },
+        "archive": {
+          "type": "string"
+        },
+        "files": {
+          "type": "string"
+        }
+      },
+      "required": [
+        "options",
+        "archive"
+      ]
+    }
+  },
+  {
+    "name": "Finish",
+    "description": "Complete the task.",
+    "parameters": {
+      "type": "object",
+      "properties": {
+        "give_answer": {
+          "type": "string"
+        }
+      },
+      "required": [
+        "give_answer"
+      ]
+    }
+  }
+]"""
+
+
+def load_model(checkpoint_path="checkpoints/linux_fc_sft/checkpoint_epoch_15.pt"):
+    """Load model and tokenizer."""
+    print("Loading tokenizer...")
+    sp = spm.SentencePieceProcessor()
+    sp.Load("tokenizer/HybriKo_tok.model")
+
+    print("Loading model...")
+    config = HybriKoConfig(
+        d_model=768, n_layers=12, vocab_size=32000,
+        n_heads=12, n_kv_heads=3, ff_mult=3,
+        max_seq_len=6144, dropout=0.0
+    )
+    model = HybriKoModel(config)
+    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+    model.load_state_dict(checkpoint["model_state_dict"])
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model.to(device).eval()
+    print(f"Model loaded on {device}\n")
+    return model, sp, device
+
+
+def generate(model, tokenizer, prompt, device, max_new_tokens=150):
+    """Generate response with improved stopping."""
+    input_ids = tokenizer.EncodeAsIds(prompt)
+    input_tensor = torch.tensor([input_ids], dtype=torch.long, device=device)
+    prompt_len = len(input_ids)
+
+    # Get EOS token ID
+    eos_id = tokenizer.PieceToId("<|im_end|>")
+    if eos_id == tokenizer.unk_id():
+        eos_id = None
+
+    with torch.no_grad():
+        generated = input_tensor
+        for _ in range(max_new_tokens):
+            outputs = model(generated)
+            logits = outputs["logits"] if isinstance(outputs, dict) else outputs.logits
+            next_token_logits = logits[:, -1, :]
+            next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
+            generated = torch.cat([generated, next_token], dim=1)
+
+            # Stop on EOS token
+            if eos_id and next_token.item() == eos_id:
+                break
+
+            # Stop when we have complete Action Input
+            new_tokens = generated[0, prompt_len:].tolist()
+            new_text = tokenizer.DecodeIds(new_tokens)
+
+            # Check for completion patterns
+            if "Action Input:" in new_text:
+                # Find the JSON part after Action Input
+                ai_idx = new_text.find("Action Input:")
+                after_ai = new_text[ai_idx + 13:].strip()
+                # Stop when JSON is complete (matching braces)
+                if after_ai.startswith("{"):
+                    brace_count = 0
+                    for i, c in enumerate(after_ai):
+                        if c == "{":
+                            brace_count += 1
+                        elif c == "}":
+                            brace_count -= 1
+                            if brace_count == 0:
+                                # Found complete JSON, stop
+                                break_idx = ai_idx + 13 + i + 1
+                                # Truncate to just the complete response
+                                new_tokens = generated[0, prompt_len:].tolist()
+                                return tokenizer.DecodeIds(new_tokens)[:break_idx + len(new_text) - len(new_text[ai_idx + 13:])]
+
+    new_tokens = generated[0, prompt_len:].tolist()
+    return tokenizer.DecodeIds(new_tokens)
+
+
+def create_prompt(user_input):
+    """Create ChatML format prompt."""
+    return f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n<|im_start|>user\n{user_input}<|im_end|>\n<|im_start|>assistant\n"
+
+
+def parse_response(response):
+    """Parse response into components."""
+    # Clean up the response - stop at <|im_end|> or garbage
+    if "<|im_end|>" in response:
+        response = response.split("<|im_end|>")[0]
+
+    # Also stop at <|im_start|> which indicates model continuing incorrectly
+    if "<|im_start|>" in response:
+        response = response.split("<|im_start|>")[0]
+
+    result = {"thought": None, "action": None, "action_input": None, "raw": response}
+
+    # Extract Thought
+    thought_match = re.search(r"Thought:\s*(.+?)(?=\s*Action:|\s*$)", response, re.DOTALL)
+    if thought_match:
+        result["thought"] = thought_match.group(1).strip()
+
+    # Extract Action
+    action_match = re.search(r"Action:\s*(\w+)", response)
+    if action_match:
+        result["action"] = action_match.group(1)
+
+    # Extract Action Input
+    input_match = re.search(r"Action Input:\s*(\{[^}]+\})", response, re.DOTALL)
+    if input_match:
+        try:
+            result["action_input"] = json.loads(input_match.group(1))
+        except:
+            result["action_input"] = input_match.group(1)
+
+    return result
+
+
+def run_single(model, tokenizer, device, user_input):
+    """Run single inference."""
+    prompt = create_prompt(user_input)
+    response = generate(model, tokenizer, prompt, device)
+    return parse_response(response)
+
+
+def main():
+    import locale
+    import io
+
+    # Set UTF-8 encoding for stdin/stdout
+    if sys.stdin.encoding != 'utf-8':
+        sys.stdin = io.TextIOWrapper(sys.stdin.buffer, encoding='utf-8', errors='replace')
+    if sys.stdout.encoding != 'utf-8':
+        sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace')
+
+    parser = argparse.ArgumentParser(description="HybriKo Linux FC Demo")
+    parser.add_argument("--checkpoint", default="checkpoints/linux_fc_sft/checkpoint_epoch_15.pt")
+    parser.add_argument("--query", type=str, help="Single query mode (non-interactive)")
+    args = parser.parse_args()
+
+    print("=" * 60)
+    print("  HybriKo-117M Linux Function Calling Demo")
+    print("=" * 60)
+
+    model, tokenizer, device = load_model(args.checkpoint)
+
+    # Single query mode
+    if args.query:
+        result = run_single(model, tokenizer, device, args.query)
+        print(f"Input: {args.query}")
+        print("-" * 40)
+        if result["thought"]:
+            print(f"Thought: {result['thought']}")
+        if result["action"]:
+            print(f"Action:  {result['action']}")
+        if result["action_input"]:
+            print(f"Input:   {json.dumps(result['action_input'], ensure_ascii=False)}")
+        return
+
+    # Interactive mode
+    print("Supported commands:")
+    print("  ls, cd, mkdir, rm, cp, mv, find, cat, grep, head,")
+    print("  tail, wc, ps, df, du, top, ping, curl, chmod, tar")
+    print("\nType 'quit' or 'exit' to exit")
+    print("=" * 60)
+
+    while True:
+        try:
+            print("\n[User] ", end="", flush=True)
+            user_input = sys.stdin.readline()
+            if not user_input:  # EOF
+                break
+            user_input = user_input.strip()
+            if not user_input:
+                continue
+            if user_input.lower() in ["quit", "exit", "q"]:
+                print("Goodbye!")
+                break
+
+            result = run_single(model, tokenizer, device, user_input)
+
+            print("\n[HybriKo]")
+            print("-" * 40)
+            if result["thought"]:
+                print(f"Thought: {result['thought']}")
+            if result["action"]:
+                print(f"Action:  {result['action']}")
+            if result["action_input"]:
+                print(f"Input:   {json.dumps(result['action_input'], ensure_ascii=False)}")
+            print("-" * 40)
+
+        except KeyboardInterrupt:
+            print("\nGoodbye!")
+            break
+        except UnicodeDecodeError as e:
+            print(f"\n[Error] Encoding issue: {e}")
+            print("Try using --query option instead: python scripts/demo_linux_fc.py --query \"your query\"")
+            continue
+
+
+if __name__ == "__main__":
+    main()