Create yen.py
Browse files
yen.py
ADDED
|
@@ -0,0 +1,17 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import torch, torch.nn as nn
|
| 2 |
+
from huggingface_hub import hf_hub_download
|
| 3 |
+
from safetensors.torch import load_file
|
| 4 |
+
|
| 5 |
+
class OpenAI:
|
| 6 |
+
def __init__(self, token=None):
|
| 7 |
+
self.d = "cuda" if torch.cuda.is_available() else "cpu"
|
| 8 |
+
class M(nn.Module):
|
| 9 |
+
def __init__(self):
|
| 10 |
+
super().__init__(); self.embedding = nn.Embedding(256, 1); self.classifier = nn.Linear(1, 2, bias=False)
|
| 11 |
+
def forward(self, x): return self.classifier(self.embedding(x).float().mean(0))
|
| 12 |
+
self.m = M().to(self.d); self.m.load_state_dict(load_file(hf_hub_download("faunix/YEN", "model.safetensors", token=token)))
|
| 13 |
+
self.chat = type('o', (object,), {"completions": type('o', (object,), {"create": self._c})()})()
|
| 14 |
+
def _c(self, **k):
|
| 15 |
+
t = torch.tensor([ord(c)%256 for c in k['messages'][-1]['content']], device=self.d)
|
| 16 |
+
with torch.no_grad(): out = self.m(t); res = "Yes" if out[1] > out[0] else "No"
|
| 17 |
+
return type('o', (object,), {"choices": [type('o', (object,), {"message": type('o', (object,), {"content": res})()})()]})()
|