初始化项目，由ModelHub XC社区提供模型
Model: LoganResearch/ARC-Base-8B Source: Original Platform
This commit is contained in:
ModelHub XC
631
inference.py
Normal file
631
inference.py
Normal file
@@ -0,0 +1,631 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
ARC-8B: Adaptive Repetition Controller
|
||||
=======================================
|
||||
Decode-time behavioral control for language models.
|
||||
|
||||
This script loads the complete ARC system and runs inference with
|
||||
multi-head cognitive control that detects and suppresses:
|
||||
- Repetition loops (125× separation)
|
||||
- Hedging phrases (1.5× separation)
|
||||
- Verbosity/filler (2.1× separation)
|
||||
- Sycophancy (experimental)
|
||||
|
||||
Usage:
|
||||
python inference.py # Interactive mode
|
||||
python inference.py --prompt "Hello" # Single prompt
|
||||
python inference.py --no-arc # Disable ARC (baseline)
|
||||
|
||||
Requirements:
|
||||
pip install torch transformers accelerate bitsandbytes
|
||||
|
||||
Model: LoganResearch/ARC-Base-8B (16GB, runs in ~10GB with 4-bit)
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import argparse
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# CONFIGURATION
|
||||
# =============================================================================
|
||||
|
||||
@dataclass
|
||||
class ARCConfig:
|
||||
"""ARC System Configuration"""
|
||||
# Model
|
||||
model_id: str = "LoganResearch/ARC-Base-8B"
|
||||
load_in_4bit: bool = True
|
||||
load_in_8bit: bool = False
|
||||
device_map: str = "auto"
|
||||
|
||||
# Architecture (must match training)
|
||||
d_model: int = 4096
|
||||
n_layers: int = 32
|
||||
d_fiber: int = 16
|
||||
d_control: int = 64
|
||||
|
||||
# Intervention thresholds (tuned empirically)
|
||||
repetition_threshold: float = 0.70
|
||||
hedging_threshold: float = 0.60
|
||||
verbosity_threshold: float = 0.65
|
||||
sycophancy_threshold: float = 0.60
|
||||
|
||||
# Intervention penalties
|
||||
repetition_penalty: float = 5.0
|
||||
hedging_penalty: float = 3.0
|
||||
verbosity_penalty: float = 2.0
|
||||
sycophancy_penalty: float = 2.0
|
||||
|
||||
# Generation
|
||||
max_new_tokens: int = 512
|
||||
temperature: float = 0.8
|
||||
top_p: float = 0.92
|
||||
repetition_window: int = 32
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# MULTI-HEAD PREDICTOR
|
||||
# =============================================================================
|
||||
|
||||
class MultiHeadPredictor(nn.Module):
|
||||
"""
|
||||
Prediction heads that monitor hidden states and detect behavioral patterns.
|
||||
|
||||
The system uses shared "fiber projections" that compress hidden states,
|
||||
then individual heads that predict risk scores for specific behaviors.
|
||||
|
||||
Architecture:
|
||||
Hidden States [n_layers × d_model]
|
||||
→ Fiber Projections [n_layers × d_fiber]
|
||||
→ Weighted Aggregation [d_fiber]
|
||||
→ Per-Head MLP → Risk Score [0-1]
|
||||
"""
|
||||
|
||||
def __init__(self, config: ARCConfig):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
|
||||
# Shared fiber projections (learned during CF-HoT training)
|
||||
self.fiber_projs = nn.ModuleList([
|
||||
nn.Linear(config.d_model, config.d_fiber, bias=False)
|
||||
for _ in range(config.n_layers)
|
||||
])
|
||||
|
||||
# Learned layer importance weights
|
||||
self.layer_weights = nn.Parameter(torch.ones(config.n_layers) / config.n_layers)
|
||||
|
||||
# Individual prediction heads
|
||||
self.heads = nn.ModuleDict()
|
||||
self.loaded_heads: set = set()
|
||||
|
||||
def _make_head(self) -> nn.Sequential:
|
||||
"""Create a prediction head: fiber features → risk score"""
|
||||
return nn.Sequential(
|
||||
nn.Linear(self.config.d_fiber, self.config.d_control),
|
||||
nn.GELU(),
|
||||
nn.Linear(self.config.d_control, self.config.d_control),
|
||||
nn.GELU(),
|
||||
nn.Linear(self.config.d_control, 1)
|
||||
)
|
||||
|
||||
def add_head(self, name: str) -> None:
|
||||
"""Add a new prediction head"""
|
||||
self.heads[name] = self._make_head()
|
||||
|
||||
def get_fiber_features(self, hidden_states: List[torch.Tensor]) -> torch.Tensor:
|
||||
"""
|
||||
Project hidden states through fiber projections and aggregate.
|
||||
|
||||
Args:
|
||||
hidden_states: List of [batch, seq, d_model] tensors from each layer
|
||||
|
||||
Returns:
|
||||
Aggregated features [batch, seq, d_fiber]
|
||||
"""
|
||||
device = hidden_states[0].device
|
||||
fibers = []
|
||||
for i, (proj, hidden) in enumerate(zip(self.fiber_projs, hidden_states)):
|
||||
if i < len(hidden_states):
|
||||
proj = proj.to(device)
|
||||
fibers.append(proj(hidden.float()))
|
||||
|
||||
# Weighted sum across layers
|
||||
weights = F.softmax(self.layer_weights.to(device)[:len(fibers)], dim=0)
|
||||
aggregated = sum(w * f for w, f in zip(weights, fibers))
|
||||
return aggregated
|
||||
|
||||
def get_risk(self, head_name: str, hidden_states: List[torch.Tensor]) -> torch.Tensor:
|
||||
"""Get risk score from a specific head"""
|
||||
if head_name not in self.loaded_heads:
|
||||
return torch.zeros(1, device=hidden_states[0].device)
|
||||
|
||||
features = self.get_fiber_features(hidden_states)
|
||||
logits = self.heads[head_name](features).squeeze(-1)
|
||||
return torch.sigmoid(logits)
|
||||
|
||||
def get_all_risks(self, hidden_states: List[torch.Tensor]) -> Dict[str, torch.Tensor]:
|
||||
"""Get risk scores from all loaded heads"""
|
||||
if not self.loaded_heads:
|
||||
return {}
|
||||
|
||||
device = hidden_states[0].device
|
||||
features = self.get_fiber_features(hidden_states)
|
||||
risks = {}
|
||||
for name in self.loaded_heads:
|
||||
self.heads[name] = self.heads[name].to(device)
|
||||
logits = self.heads[name](features).squeeze(-1)
|
||||
risks[name] = torch.sigmoid(logits)
|
||||
return risks
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# ARC SYSTEM
|
||||
# =============================================================================
|
||||
|
||||
class ARCSystem:
|
||||
"""
|
||||
Complete ARC (Adaptive Repetition Controller) System
|
||||
|
||||
Loads model + prediction heads and provides controlled generation
|
||||
with real-time behavioral intervention.
|
||||
"""
|
||||
|
||||
# Tokens to suppress for each behavior type
|
||||
HEDGE_STARTERS = [
|
||||
"As", "I'm", "I", "It's", "While", "Although", "However",
|
||||
"That", "This", "Please", "Well", "So", "Actually"
|
||||
]
|
||||
VERBOSE_STARTERS = [
|
||||
"Let", "Basically", "Essentially", "Simply", "Indeed",
|
||||
"Furthermore", "Moreover", "Additionally", "Firstly"
|
||||
]
|
||||
SYCOPHANCY_STARTERS = [
|
||||
"Great", "Excellent", "Wonderful", "Absolutely", "Of",
|
||||
"Thank", "Sure", "Certainly", "Definitely"
|
||||
]
|
||||
|
||||
def __init__(self, config: Optional[ARCConfig] = None):
|
||||
self.config = config or ARCConfig()
|
||||
|
||||
self.model = None
|
||||
self.tokenizer = None
|
||||
self.predictor = None
|
||||
|
||||
# Token ID caches for suppression
|
||||
self._hedge_token_ids: set = set()
|
||||
self._verbose_token_ids: set = set()
|
||||
self._sycophancy_token_ids: set = set()
|
||||
|
||||
# Stats
|
||||
self.total_interventions = {"repetition": 0, "hedging": 0, "verbosity": 0, "sycophancy": 0}
|
||||
|
||||
def load(self, verbose: bool = True) -> "ARCSystem":
|
||||
"""
|
||||
Load all components from HuggingFace.
|
||||
|
||||
Downloads and initializes:
|
||||
1. Base model (Hermes-3-Llama-3.1-8B based)
|
||||
2. Tokenizer
|
||||
3. Prediction heads (repetition, hedging, verbosity, sycophancy)
|
||||
|
||||
Returns:
|
||||
self (for chaining)
|
||||
"""
|
||||
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
if verbose:
|
||||
print("=" * 60)
|
||||
print(" ARC-8B: Adaptive Repetition Controller")
|
||||
print(" Decode-time behavioral control system")
|
||||
print("=" * 60)
|
||||
|
||||
# === 1. Tokenizer ===
|
||||
if verbose:
|
||||
print("\n[1/4] Loading tokenizer...")
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(
|
||||
self.config.model_id,
|
||||
trust_remote_code=True
|
||||
)
|
||||
if self.tokenizer.pad_token is None:
|
||||
self.tokenizer.pad_token = self.tokenizer.eos_token
|
||||
|
||||
# === 2. Model ===
|
||||
if verbose:
|
||||
print("[2/4] Loading model...")
|
||||
if self.config.load_in_4bit:
|
||||
print(" (4-bit quantization enabled)")
|
||||
|
||||
quantization_config = None
|
||||
if self.config.load_in_4bit:
|
||||
quantization_config = BitsAndBytesConfig(
|
||||
load_in_4bit=True,
|
||||
bnb_4bit_compute_dtype=torch.float16,
|
||||
bnb_4bit_use_double_quant=True,
|
||||
bnb_4bit_quant_type="nf4"
|
||||
)
|
||||
elif self.config.load_in_8bit:
|
||||
quantization_config = BitsAndBytesConfig(load_in_8bit=True)
|
||||
|
||||
self.model = AutoModelForCausalLM.from_pretrained(
|
||||
self.config.model_id,
|
||||
quantization_config=quantization_config,
|
||||
device_map=self.config.device_map,
|
||||
torch_dtype=torch.float16,
|
||||
trust_remote_code=True
|
||||
)
|
||||
self.model.eval()
|
||||
|
||||
# === 3. Prediction Heads ===
|
||||
if verbose:
|
||||
print("[3/4] Loading prediction heads...")
|
||||
|
||||
device = next(self.model.parameters()).device
|
||||
self.predictor = MultiHeadPredictor(self.config).to(device).float()
|
||||
|
||||
# Load risk_predictor.pt (contains fiber projections + repetition head)
|
||||
try:
|
||||
risk_path = hf_hub_download(self.config.model_id, "risk_predictor.pt")
|
||||
ckpt = torch.load(risk_path, map_location=device, weights_only=False)
|
||||
|
||||
# The checkpoint contains the full state dict
|
||||
state = ckpt.get('risk_predictor', ckpt)
|
||||
|
||||
# Load fiber projections
|
||||
for i in range(self.config.n_layers):
|
||||
key = f'fiber_projs.{i}.weight'
|
||||
if key in state:
|
||||
self.predictor.fiber_projs[i].weight.data = state[key].to(device).float()
|
||||
|
||||
# Load layer weights
|
||||
if 'layer_weights' in state:
|
||||
self.predictor.layer_weights.data = state['layer_weights'].to(device).float()
|
||||
|
||||
# Load repetition head
|
||||
self.predictor.add_head('repetition')
|
||||
self.predictor.heads['repetition'][0].weight.data = state['predictor.0.weight'].to(device).float()
|
||||
self.predictor.heads['repetition'][0].bias.data = state['predictor.0.bias'].to(device).float()
|
||||
self.predictor.heads['repetition'][2].weight.data = state['predictor.2.weight'].to(device).float()
|
||||
self.predictor.heads['repetition'][2].bias.data = state['predictor.2.bias'].to(device).float()
|
||||
self.predictor.heads['repetition'][4].weight.data = state['predictor.4.weight'].to(device).float()
|
||||
self.predictor.heads['repetition'][4].bias.data = state['predictor.4.bias'].to(device).float()
|
||||
self.predictor.loaded_heads.add('repetition')
|
||||
|
||||
if verbose:
|
||||
print(" ✓ Repetition head (125× separation)")
|
||||
except Exception as e:
|
||||
if verbose:
|
||||
print(f" ✗ Repetition head: {e}")
|
||||
|
||||
# Load additional heads
|
||||
for head_name in ['hedging', 'verbosity', 'sycophancy']:
|
||||
try:
|
||||
head_path = hf_hub_download(self.config.model_id, f"{head_name}_head.pt")
|
||||
ckpt = torch.load(head_path, map_location=device, weights_only=False)
|
||||
|
||||
self.predictor.add_head(head_name)
|
||||
head_state = ckpt.get('head_state', ckpt)
|
||||
self.predictor.heads[head_name].load_state_dict(head_state)
|
||||
self.predictor.loaded_heads.add(head_name)
|
||||
|
||||
if verbose:
|
||||
print(f" ✓ {head_name.capitalize()} head")
|
||||
except Exception as e:
|
||||
if verbose:
|
||||
print(f" ✗ {head_name.capitalize()} head: {e}")
|
||||
|
||||
self.predictor.eval()
|
||||
|
||||
# === 4. Build Token Suppression Sets ===
|
||||
if verbose:
|
||||
print("[4/4] Building suppression vocabularies...")
|
||||
|
||||
self._build_suppression_sets()
|
||||
|
||||
if verbose:
|
||||
print("\n" + "=" * 60)
|
||||
print(f" ✓ ARC System Ready")
|
||||
print(f" Active heads: {list(self.predictor.loaded_heads)}")
|
||||
print("=" * 60 + "\n")
|
||||
|
||||
return self
|
||||
|
||||
def _build_suppression_sets(self) -> None:
|
||||
"""Build token ID sets for behavioral suppression"""
|
||||
for word in self.HEDGE_STARTERS:
|
||||
tokens = self.tokenizer.encode(word, add_special_tokens=False)
|
||||
if tokens:
|
||||
self._hedge_token_ids.add(tokens[0])
|
||||
|
||||
for word in self.VERBOSE_STARTERS:
|
||||
tokens = self.tokenizer.encode(word, add_special_tokens=False)
|
||||
if tokens:
|
||||
self._verbose_token_ids.add(tokens[0])
|
||||
|
||||
for word in self.SYCOPHANCY_STARTERS:
|
||||
tokens = self.tokenizer.encode(word, add_special_tokens=False)
|
||||
if tokens:
|
||||
self._sycophancy_token_ids.add(tokens[0])
|
||||
|
||||
def _apply_interventions(
|
||||
self,
|
||||
logits: torch.Tensor,
|
||||
risks: Dict[str, torch.Tensor],
|
||||
recent_tokens: List[int]
|
||||
) -> Tuple[torch.Tensor, Dict[str, bool]]:
|
||||
"""
|
||||
Apply behavioral interventions based on risk scores.
|
||||
|
||||
Args:
|
||||
logits: [1, vocab_size] logits for next token
|
||||
risks: Dict of risk scores for each head
|
||||
recent_tokens: Recently generated token IDs
|
||||
|
||||
Returns:
|
||||
Modified logits and dict of which interventions fired
|
||||
"""
|
||||
interventions = {}
|
||||
|
||||
# Repetition: suppress recently used tokens
|
||||
if risks.get('repetition', 0) > self.config.repetition_threshold:
|
||||
for tok in set(recent_tokens[-self.config.repetition_window:]):
|
||||
logits[0, tok] -= self.config.repetition_penalty
|
||||
interventions['repetition'] = True
|
||||
self.total_interventions['repetition'] += 1
|
||||
|
||||
# Hedging: suppress hedge phrase starters
|
||||
if risks.get('hedging', 0) > self.config.hedging_threshold:
|
||||
for tok in self._hedge_token_ids:
|
||||
logits[0, tok] -= self.config.hedging_penalty
|
||||
interventions['hedging'] = True
|
||||
self.total_interventions['hedging'] += 1
|
||||
|
||||
# Verbosity: suppress filler phrase starters
|
||||
if risks.get('verbosity', 0) > self.config.verbosity_threshold:
|
||||
for tok in self._verbose_token_ids:
|
||||
logits[0, tok] -= self.config.verbosity_penalty
|
||||
interventions['verbosity'] = True
|
||||
self.total_interventions['verbosity'] += 1
|
||||
|
||||
# Sycophancy: suppress sycophantic starters
|
||||
if risks.get('sycophancy', 0) > self.config.sycophancy_threshold:
|
||||
for tok in self._sycophancy_token_ids:
|
||||
logits[0, tok] -= self.config.sycophancy_penalty
|
||||
interventions['sycophancy'] = True
|
||||
self.total_interventions['sycophancy'] += 1
|
||||
|
||||
return logits, interventions
|
||||
|
||||
def generate(
|
||||
self,
|
||||
prompt: str,
|
||||
system_prompt: Optional[str] = None,
|
||||
max_new_tokens: Optional[int] = None,
|
||||
temperature: Optional[float] = None,
|
||||
use_arc: bool = True,
|
||||
verbose: bool = False
|
||||
) -> str:
|
||||
"""
|
||||
Generate text with optional ARC behavioral control.
|
||||
|
||||
Args:
|
||||
prompt: User input
|
||||
system_prompt: Optional system message
|
||||
max_new_tokens: Max tokens to generate (default: config value)
|
||||
temperature: Sampling temperature (default: config value)
|
||||
use_arc: Whether to use ARC intervention (default: True)
|
||||
verbose: Print intervention info (default: False)
|
||||
|
||||
Returns:
|
||||
Generated text
|
||||
"""
|
||||
max_new_tokens = max_new_tokens or self.config.max_new_tokens
|
||||
temperature = temperature or self.config.temperature
|
||||
|
||||
# Build chat format
|
||||
if system_prompt is None:
|
||||
system_prompt = "You are a helpful assistant."
|
||||
|
||||
full_prompt = f"<|im_start|>system\n{system_prompt}<|im_end|>\n"
|
||||
full_prompt += f"<|im_start|>user\n{prompt}<|im_end|>\n"
|
||||
full_prompt += "<|im_start|>assistant\n"
|
||||
|
||||
device = next(self.model.parameters()).device
|
||||
input_ids = self.tokenizer.encode(full_prompt, return_tensors='pt').to(device)
|
||||
attention_mask = torch.ones_like(input_ids)
|
||||
|
||||
generated_ids = input_ids.clone()
|
||||
intervention_counts = {"repetition": 0, "hedging": 0, "verbosity": 0, "sycophancy": 0}
|
||||
|
||||
# Generation loop
|
||||
for step in range(max_new_tokens):
|
||||
with torch.no_grad():
|
||||
outputs = self.model(
|
||||
input_ids=generated_ids,
|
||||
attention_mask=attention_mask,
|
||||
output_hidden_states=True,
|
||||
return_dict=True
|
||||
)
|
||||
|
||||
logits = outputs.logits[:, -1, :] / temperature
|
||||
|
||||
# ARC intervention
|
||||
if use_arc and self.predictor.loaded_heads:
|
||||
hidden_states = outputs.hidden_states[1:] # Skip embedding layer
|
||||
risks = self.predictor.get_all_risks(hidden_states)
|
||||
current_risks = {name: r[:, -1].item() for name, r in risks.items()}
|
||||
|
||||
recent = generated_ids[0, -self.config.repetition_window:].tolist()
|
||||
logits, fired = self._apply_interventions(logits, current_risks, recent)
|
||||
|
||||
for k, v in fired.items():
|
||||
if v:
|
||||
intervention_counts[k] += 1
|
||||
|
||||
# Top-p sampling
|
||||
sorted_logits, sorted_indices = torch.sort(logits, descending=True)
|
||||
cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
|
||||
sorted_indices_to_remove = cumulative_probs > self.config.top_p
|
||||
sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
|
||||
sorted_indices_to_remove[..., 0] = 0
|
||||
indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
|
||||
logits[indices_to_remove] = float('-inf')
|
||||
|
||||
probs = F.softmax(logits, dim=-1)
|
||||
next_token = torch.multinomial(probs, num_samples=1)
|
||||
|
||||
generated_ids = torch.cat([generated_ids, next_token], dim=-1)
|
||||
attention_mask = torch.cat([attention_mask, torch.ones(1, 1, device=device)], dim=-1)
|
||||
|
||||
# Check for EOS
|
||||
if next_token.item() == self.tokenizer.eos_token_id:
|
||||
break
|
||||
|
||||
# Check for end of turn
|
||||
if next_token.item() == self.tokenizer.encode("<|im_end|>", add_special_tokens=False)[0]:
|
||||
break
|
||||
|
||||
# Decode response
|
||||
full_output = self.tokenizer.decode(generated_ids[0], skip_special_tokens=False)
|
||||
|
||||
# Extract assistant response
|
||||
if "<|im_start|>assistant\n" in full_output:
|
||||
response = full_output.split("<|im_start|>assistant\n")[-1]
|
||||
if "<|im_end|>" in response:
|
||||
response = response.split("<|im_end|>")[0]
|
||||
else:
|
||||
response = full_output
|
||||
|
||||
if verbose:
|
||||
total = sum(intervention_counts.values())
|
||||
print(f"\n[ARC Stats] Interventions: {total} total")
|
||||
for k, v in intervention_counts.items():
|
||||
if v > 0:
|
||||
print(f" - {k}: {v}")
|
||||
|
||||
return response.strip()
|
||||
|
||||
def chat(self, system_prompt: Optional[str] = None) -> None:
|
||||
"""
|
||||
Interactive chat mode.
|
||||
|
||||
Args:
|
||||
system_prompt: Optional system message
|
||||
"""
|
||||
print("\n" + "=" * 60)
|
||||
print(" ARC-8B Interactive Chat")
|
||||
print(" Commands: /quit, /stats, /arc on|off, /clear")
|
||||
print("=" * 60 + "\n")
|
||||
|
||||
use_arc = True
|
||||
history = []
|
||||
|
||||
while True:
|
||||
try:
|
||||
user_input = input("You: ").strip()
|
||||
except (KeyboardInterrupt, EOFError):
|
||||
print("\nGoodbye!")
|
||||
break
|
||||
|
||||
if not user_input:
|
||||
continue
|
||||
|
||||
# Commands
|
||||
if user_input.lower() == '/quit':
|
||||
print("Goodbye!")
|
||||
break
|
||||
elif user_input.lower() == '/stats':
|
||||
print(f"\nTotal interventions: {self.total_interventions}\n")
|
||||
continue
|
||||
elif user_input.lower() == '/arc on':
|
||||
use_arc = True
|
||||
print("ARC enabled\n")
|
||||
continue
|
||||
elif user_input.lower() == '/arc off':
|
||||
use_arc = False
|
||||
print("ARC disabled (baseline mode)\n")
|
||||
continue
|
||||
elif user_input.lower() == '/clear':
|
||||
history = []
|
||||
self.total_interventions = {k: 0 for k in self.total_interventions}
|
||||
print("History cleared\n")
|
||||
continue
|
||||
|
||||
# Generate response
|
||||
response = self.generate(
|
||||
user_input,
|
||||
system_prompt=system_prompt,
|
||||
use_arc=use_arc,
|
||||
verbose=True
|
||||
)
|
||||
|
||||
print(f"\nAssistant: {response}\n")
|
||||
history.append({"user": user_input, "assistant": response})
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# MAIN
|
||||
# =============================================================================
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="ARC-8B: Adaptive Repetition Controller",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog="""
|
||||
Examples:
|
||||
python inference.py # Interactive chat
|
||||
python inference.py --prompt "Hello" # Single prompt
|
||||
python inference.py --no-arc # Disable ARC (baseline)
|
||||
python inference.py --8bit # Use 8-bit quantization
|
||||
"""
|
||||
)
|
||||
parser.add_argument("--prompt", "-p", type=str, help="Single prompt to process")
|
||||
parser.add_argument("--system", "-s", type=str, help="System prompt")
|
||||
parser.add_argument("--no-arc", action="store_true", help="Disable ARC intervention")
|
||||
parser.add_argument("--4bit", dest="load_4bit", action="store_true", default=True, help="Use 4-bit quantization (default)")
|
||||
parser.add_argument("--8bit", dest="load_8bit", action="store_true", help="Use 8-bit quantization")
|
||||
parser.add_argument("--no-quant", action="store_true", help="Disable quantization (requires ~32GB VRAM)")
|
||||
parser.add_argument("--max-tokens", type=int, default=512, help="Max tokens to generate")
|
||||
parser.add_argument("--temperature", type=float, default=0.8, help="Sampling temperature")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Configure
|
||||
config = ARCConfig(
|
||||
max_new_tokens=args.max_tokens,
|
||||
temperature=args.temperature
|
||||
)
|
||||
|
||||
if args.load_8bit:
|
||||
config.load_in_4bit = False
|
||||
config.load_in_8bit = True
|
||||
elif args.no_quant:
|
||||
config.load_in_4bit = False
|
||||
config.load_in_8bit = False
|
||||
|
||||
# Load
|
||||
arc = ARCSystem(config)
|
||||
arc.load()
|
||||
|
||||
# Run
|
||||
if args.prompt:
|
||||
response = arc.generate(
|
||||
args.prompt,
|
||||
system_prompt=args.system,
|
||||
use_arc=not args.no_arc,
|
||||
verbose=True
|
||||
)
|
||||
print(f"\n{response}\n")
|
||||
else:
|
||||
arc.chat(system_prompt=args.system)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user