Limen4ik commited on
Commit
a46c7cd
·
verified ·
1 Parent(s): 1559b9c

Create yen.py

Browse files
Files changed (1) hide show
  1. yen.py +17 -0
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})()})()]})()