alex-abb's picture
Update app.py
16fd66a verified
raw
history blame
1.79 kB
import datasets
from datasets import load_dataset
import transformers
from transformers import AutoModelForSequenceClassification, AutoTokenizer, Trainer, TrainingArguments, pipeline
model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2) # 2 classes : positif et négatif
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
ds = load_dataset("stanfordnlp/sst2")
sst2_dataset = load_dataset("glue", "sst2", split="train")
def encode(examples):
return tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, padding="max_length")
sst2_dataset = sst2_dataset.map(encode, batched=True)
sst2_dataset = sst2_dataset.map(lambda examples: {"labels": examples["label"]}, batched=True)
training_args = TrainingArguments(
per_device_train_batch_size=8,
evaluation_strategy="epoch",
logging_dir="./logs",
output_dir="./results",
num_train_epochs=3,
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=encoded_dataset["train"],
eval_dataset=encoded_dataset["test"],
)
import os
if not os.path.exists("./fine_tuned_model"):
trainer.train()
# Sauvegarder le modèle fine-tuné et le tokenizer
model.save_pretrained("./fine_tuned_model")
tokenizer.save_pretrained("./fine_tuned_model")
else:
# Charger le modèle fine-tuné
model = BertForSequenceClassification.from_pretrained("./fine_tuned_model")
tokenizer = BertTokenizer.from_pretrained("./fine_tuned_model")
sentiment_analysis = pipeline("sentiment-analysis", model=model, tokenizer=tokenizer)
def generate_response(message):
result = sentiment_analysis(message)[0]
return f"Label: {result['label']}, Score: {result['score']}"
gr.ChatInterface(fn=generate_response).launch()