Tbruand commited on
Commit
88ca48d
·
1 Parent(s): 3e3daf9

feat(handler): ajoute la prise en charge du modèle fine-tuné dans la prédiction

Browse files
Files changed (1) hide show
  1. app/handler.py +12 -8
app/handler.py CHANGED
@@ -1,19 +1,23 @@
1
  from models.zero_shot import ZeroShotModel
2
  from models.few_shot import FewShotModel
 
3
 
4
  zero_shot_model = ZeroShotModel()
5
  few_shot_model = FewShotModel()
 
6
 
7
  def predict(text: str, model_type: str = "zero-shot") -> str:
8
  if model_type == "few-shot":
9
  results = few_shot_model.predict(text)
10
- output = "### Résultat de la classification (Few-Shot) :\n\n"
11
- for label, score in results:
12
- output += f"- **{label}** : {score*100:.1f}%\n"
13
- return output
14
  else:
15
  results = zero_shot_model.predict(text)
16
- output = "### Résultat de la classification (Zero-Shot) :\n\n"
17
- for label, score in results:
18
- output += f"- **{label}** : {score*100:.1f}%\n"
19
- return output
 
 
 
1
  from models.zero_shot import ZeroShotModel
2
  from models.few_shot import FewShotModel
3
+ from models.fine_tuned import FineTunedModel
4
 
5
  zero_shot_model = ZeroShotModel()
6
  few_shot_model = FewShotModel()
7
+ fine_tuned_model = FineTunedModel()
8
 
9
  def predict(text: str, model_type: str = "zero-shot") -> str:
10
  if model_type == "few-shot":
11
  results = few_shot_model.predict(text)
12
+ title = "Few-Shot"
13
+ elif model_type == "fine-tuned":
14
+ results = fine_tuned_model.predict(text)
15
+ title = "Fine-Tuned"
16
  else:
17
  results = zero_shot_model.predict(text)
18
+ title = "Zero-Shot"
19
+
20
+ output = f"### Résultat de la classification ({title}) :\n\n"
21
+ for label, score in results:
22
+ output += f"- **{label}** : {score*100:.1f}%\n"
23
+ return output