104 lines
3.0 KiB
Python
Executable File
104 lines
3.0 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
Inference script for trained LoRA adapter.
|
|
|
|
Loads the trained model and generates Cyron summaries.
|
|
"""
|
|
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
|
|
def load_model(model_path):
|
|
"""Load trained LoRA model."""
|
|
from peft import PeftModel
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
|
print(f"Loading base model from {model_path}...")
|
|
base_model = AutoModelForCausalLM.from_pretrained(
|
|
model_path,
|
|
device_map="auto",
|
|
torch_dtype="auto",
|
|
)
|
|
|
|
print(f"Loading LoRA weights from {model_path}...")
|
|
model = PeftModel.from_pretrained(base_model, model_path)
|
|
tokenizer = AutoTokenizer.from_pretrained(model_path)
|
|
|
|
return model, tokenizer
|
|
|
|
|
|
def generate_summary(model, tokenizer, task, files_changed=None, tests_run=False, test_count=0, commit=False, push=False, errors_seen=False):
|
|
"""Generate a Cyron summary for a task."""
|
|
|
|
# Build prompt
|
|
prompt = f"Task: {task}"
|
|
if files_changed:
|
|
prompt += f"\nFiles: {', '.join(files_changed[:3])}"
|
|
|
|
if tests_run:
|
|
prompt += f"\nTests run: {test_count}"
|
|
if commit:
|
|
prompt += "\nCommit: yes"
|
|
if push:
|
|
prompt += "\nPush: yes"
|
|
if errors_seen:
|
|
prompt += "\nErrors seen: yes"
|
|
|
|
prompt += "\n\nGenerate a Cyron summary:"
|
|
|
|
# Tokenize
|
|
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
|
|
|
|
# Generate
|
|
outputs = model.generate(
|
|
**inputs,
|
|
max_new_tokens=500,
|
|
do_sample=True,
|
|
temperature=0.7,
|
|
top_p=0.9,
|
|
repetition_penalty=1.1,
|
|
)
|
|
|
|
# Decode
|
|
generated = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
|
|
|
|
return generated
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Generate Cyron summaries with LoRA")
|
|
parser.add_argument("--model", type=str, default="output/llama2-7b-lora",
|
|
help="Path to trained model")
|
|
parser.add_argument("--task", type=str, required=True, help="Task description")
|
|
parser.add_argument("--files", type=str, nargs="*", default=[], help="Files changed")
|
|
parser.add_argument("--tests-run", action="store_true", help="Tests were run")
|
|
parser.add_argument("--test-count", type=int, default=0, help="Number of tests")
|
|
parser.add_argument("--commit", action="store_true", help="Commit was made")
|
|
parser.add_argument("--push", action="store_true", help="Code was pushed")
|
|
parser.add_argument("--errors", action="store_true", help="Errors were seen")
|
|
args = parser.parse_args()
|
|
|
|
model, tokenizer = load_model(args.model)
|
|
|
|
summary = generate_summary(
|
|
model,
|
|
tokenizer,
|
|
task=args.task,
|
|
files_changed=args.files,
|
|
tests_run=args.tests_run,
|
|
test_count=args.test_count,
|
|
commit=args.commit,
|
|
push=args.push,
|
|
errors_seen=args.errors,
|
|
)
|
|
|
|
print("\nGenerated Summary:")
|
|
print("=" * 70)
|
|
print(summary)
|
|
print("=" * 70)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|