from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
tokenizer = AutoTokenizer.from_pretrained("zed-industries/zeta")
model = AutoModelForCausalLM.from_pretrained("zed-industries/zeta", torch_dtype=torch.bfloat16).to("cuda")
alpaca_prompt = """### Instruction:
You are a code completion assistant and your task is to analyze user edits and then rewrite an excerpt that the user provides, suggesting the appropriate edits within the excerpt, taking into account the cursor location.

### User Edits:

User edited "models/customer.rb":

```diff
@@ -2,5 +2,5 @@
class Customer
def initialize
- @name = name
+ @name = name.capitalize
@email = email
@phone = phone
```

### User Excerpt:

```models/customer.rb
def initialize
<|editable_region_start|>
@name = name.capitalize<|user_cursor_is_here|>
@email = email
@phone = phone
@address = address
end

def to_s
@name
end

<|editable_region_end|>
private

def validate_email
@email.include?('@')
```

### Response:

"""

inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_length=128)
print(tokenizer.decode(outputs[0], skip_special_tokens=True)[len(input_text):])