Adding another allowed origin. Commenting out few-shot model that is unused.

app.py

@@ -17,42 +17,42 @@ app = FastAPI()
 
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["https://statosphere-3704059fdd7e.c5v4v4jx6pq5.win","https://crunchatize-77a78ffcc6a6.c5v4v4jx6pq5.win","https://lord-raven.github.io"],
+    allow_origins=["https://statosphere-3704059fdd7e.c5v4v4jx6pq5.win","https://crunchatize-77a78ffcc6a6.c5v4v4jx6pq5.win","https://tamabotchi-2dba63df3bf1.c5v4v4jx6pq5.win","https://lord-raven.github.io"],
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-class OnnxSetFitModel:
-    def __init__(self, ort_model, tokenizer, model_head):
-        self.ort_model = ort_model
-        self.tokenizer = tokenizer
-        self.model_head = model_head
-
-    def predict(self, inputs):
-        encoded_inputs = self.tokenizer(
-            inputs, padding=True, truncation=True, return_tensors="pt"
-        ).to(self.ort_model.device)
-
-        outputs = self.ort_model(**encoded_inputs)
-        embeddings = mean_pooling(
-            outputs["last_hidden_state"], encoded_inputs["attention_mask"]
-        )
-        return self.model_head.predict(embeddings.cpu())
-
-    def predict_proba(self, inputs):
-        encoded_inputs = self.tokenizer(
-            inputs, padding=True, truncation=True, return_tensors="pt"
-        ).to(self.ort_model.device)
-
-        outputs = self.ort_model(**encoded_inputs)
-        embeddings = mean_pooling(
-            outputs["last_hidden_state"], encoded_inputs["attention_mask"]
-        )
-        return self.model_head.predict_proba(embeddings.cpu())
-
-    def __call__(self, inputs):
-        return self.predict(inputs)
+# class OnnxSetFitModel:
+#     def __init__(self, ort_model, tokenizer, model_head):
+#         self.ort_model = ort_model
+#         self.tokenizer = tokenizer
+#         self.model_head = model_head
+
+#     def predict(self, inputs):
+#         encoded_inputs = self.tokenizer(
+#             inputs, padding=True, truncation=True, return_tensors="pt"
+#         ).to(self.ort_model.device)
+
+#         outputs = self.ort_model(**encoded_inputs)
+#         embeddings = mean_pooling(
+#             outputs["last_hidden_state"], encoded_inputs["attention_mask"]
+#         )
+#         return self.model_head.predict(embeddings.cpu())
+
+#     def predict_proba(self, inputs):
+#         encoded_inputs = self.tokenizer(
+#             inputs, padding=True, truncation=True, return_tensors="pt"
+#         ).to(self.ort_model.device)
+
+#         outputs = self.ort_model(**encoded_inputs)
+#         embeddings = mean_pooling(
+#             outputs["last_hidden_state"], encoded_inputs["attention_mask"]
+#         )
+#         return self.model_head.predict_proba(embeddings.cpu())
+
+#     def __call__(self, inputs):
+#         return self.predict(inputs)
 
 # "xenova/mobilebert-uncased-mnli" "typeform/mobilebert-uncased-mnli" Fast but small--same as bundled in Statosphere
 # "xenova/deberta-v3-base-tasksource-nli" Not impressed
@@ -67,31 +67,28 @@ model = ORTModelForSequenceClassification.from_pretrained(model_name, file_name=
 tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, model_max_length=512)
 classifier = pipeline(task="zero-shot-classification", model=model, tokenizer=tokenizer)
 
-few_shot_tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-small-en-v1.5', model_max_length=512) # 'BAAI/bge-small-en-v1.5'
-ort_model = ORTModelForFeatureExtraction.from_pretrained('BAAI/bge-small-en-v1.5', file_name="onnx/model.onnx") # 'BAAI/bge-small-en-v1.5'
-few_shot_model = SetFitModel.from_pretrained("moshew/bge-small-en-v1.5_setfit-sst2-english") # "moshew/bge-small-en-v1.5_setfit-sst2-english"
+# few_shot_tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-small-en-v1.5', model_max_length=512) # 'BAAI/bge-small-en-v1.5'
+# ort_model = ORTModelForFeatureExtraction.from_pretrained('BAAI/bge-small-en-v1.5', file_name="onnx/model.onnx") # 'BAAI/bge-small-en-v1.5'
+# few_shot_model = SetFitModel.from_pretrained("moshew/bge-small-en-v1.5_setfit-sst2-english") # "moshew/bge-small-en-v1.5_setfit-sst2-english"
 
 # Train few_shot_model
-candidate_labels = ["supported", "refuted"]
-reference_dataset = load_dataset("SetFit/sst2")
-dummy_dataset = Dataset.from_dict({})
-train_dataset = get_templated_dataset(dummy_dataset, candidate_labels=candidate_labels, sample_size=8, template="The CONCLUSION is {} by the PASSAGE.")
-args = TrainingArguments(
-    batch_size=32,
-    num_epochs=1
-)
-trainer = Trainer(
-    model=few_shot_model,
-    args=args,
-    train_dataset=train_dataset,
-    eval_dataset=reference_dataset["test"]
-)
-trainer.train()
-
-# metrics = trainer.evaluate()
-# print(metrics)
-
-onnx_few_shot_model = OnnxSetFitModel(ort_model, few_shot_tokenizer, few_shot_model.model_head)
+# candidate_labels = ["supported", "refuted"]
+# reference_dataset = load_dataset("SetFit/sst2")
+# dummy_dataset = Dataset.from_dict({})
+# train_dataset = get_templated_dataset(dummy_dataset, candidate_labels=candidate_labels, sample_size=8, template="The CONCLUSION is {} by the PASSAGE.")
+# args = TrainingArguments(
+#     batch_size=32,
+#     num_epochs=1
+# )
+# trainer = Trainer(
+#     model=few_shot_model,
+#     args=args,
+#     train_dataset=train_dataset,
+#     eval_dataset=reference_dataset["test"]
+# )
+# trainer.train()
+
+# onnx_few_shot_model = OnnxSetFitModel(ort_model, few_shot_tokenizer, few_shot_model.model_head)
 
 
 
@@ -100,10 +97,10 @@ def classify(data_string, request: gradio.Request):
         if request.headers["origin"] not in ["https://statosphere-3704059fdd7e.c5v4v4jx6pq5.win", "https://crunchatize-77a78ffcc6a6.c5v4v4jx6pq5.win", "https://ravenok-statosphere-backend.hf.space", "https://lord-raven.github.io"]:
             return "{}"
     data = json.loads(data_string)
-    if 'task' in data and data['task'] == 'few_shot_classification':
-        return few_shot_classification(data)
-    else:
-        return zero_shot_classification(data)
+    # if 'task' in data and data['task'] == 'few_shot_classification':
+    #     return few_shot_classification(data)
+    # else:
+    return zero_shot_classification(data)
 
 def zero_shot_classification(data):
     results = classifier(data['sequence'], candidate_labels=data['candidate_labels'], hypothesis_template=data['hypothesis_template'], multi_label=data['multi_label'])
@@ -114,22 +111,22 @@ def create_sequences(data):
     # return ['###Given:\n' + data['sequence'] + '\n###End Given\n###Hypothesis:\n' + data['hypothesis_template'].format(label) + "\n###End Hypothesis" for label in data['candidate_labels']]
     return [data['sequence'] + '\n' + data['hypothesis_template'].format(label) for label in data['candidate_labels']]
 
-def few_shot_classification(data):
-    sequences = create_sequences(data)
-    print(sequences)
-    # results = onnx_few_shot_model(sequences)
-    probs = onnx_few_shot_model.predict_proba(sequences)
-    scores = [true[0] for true in probs]
+# def few_shot_classification(data):
+#     sequences = create_sequences(data)
+#     print(sequences)
+#     # results = onnx_few_shot_model(sequences)
+#     probs = onnx_few_shot_model.predict_proba(sequences)
+#     scores = [true[0] for true in probs]
 
-    composite = list(zip(scores, data['candidate_labels']))
-    composite = sorted(composite, key=lambda x: x[0], reverse=True)
+#     composite = list(zip(scores, data['candidate_labels']))
+#     composite = sorted(composite, key=lambda x: x[0], reverse=True)
 
-    labels, scores = zip(*composite)
+#     labels, scores = zip(*composite)
 
-    response_dict = {'scores': scores, 'labels': labels}
-    print(response_dict)
-    response_string = json.dumps(response_dict)
-    return response_string
+#     response_dict = {'scores': scores, 'labels': labels}
+#     print(response_dict)
+#     response_string = json.dumps(response_dict)
+#     return response_string
 
 gradio_interface = gradio.Interface(
     fn = classify,