File size: 1,853 Bytes
80c2763 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
# -*- coding: utf-8 -*-
"""Untitled31.ipynb
Automatically generated by Colab.
Original file is located at
https://colab.research.google.com/drive/1qkQ5UtvWMcQKdxpkEgZhStoBplz9kcAo
"""
from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
from datasets import Dataset
import gradio as gr
import torch
data = {
"text": [
"Proficient in Python and Machine Learning",
"Excellent written and verbal communication",
"Experience with cloud platforms like AWS and Azure",
"Skilled in data visualization and analytics",
"Project management and Agile methodologies"
],
"label": [0, 1, 0, 0, 1] # 0 = Technical, 1 = Soft Skill
}
dataset = Dataset.from_dict(data)
model_checkpoint = "distilbert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
def tokenize(batch):
return tokenizer(batch["text"], padding=True, truncation=True)
tokenized_dataset = dataset.map(tokenize, batched=True)
model = AutoModelForSequenceClassification.from_pretrained(model_checkpoint, num_labels=2)
training_args = TrainingArguments(
output_dir="./results",
evaluation_strategy="no",
per_device_train_batch_size=2,
num_train_epochs=3,
logging_steps=10,
push_to_hub=False,
report_to="none"
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=tokenized_dataset
)
trainer.train()
def classify(text):
inputs = tokenizer(text, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
prediction = torch.argmax(outputs.logits, dim=1).item()
return "Soft Skill" if prediction == 1 else "Technical Skill"
print(classify("Familiar with cloud computing and Docker"))
interface = gr.Interface(fn=classify, inputs="text", outputs="text")
interface.launch()
|