Spaces:
Sleeping
Sleeping
import streamlit as st | |
from transformers import DistilBertTokenizerFast, DistilBertForSequenceClassification | |
import torch | |
import torch.nn.functional as F | |
import zipfile | |
import shutil | |
import os | |
def unzip_and_save(zip_file_path, extraction_path): | |
# Create the extraction directory if it doesn't exist | |
os.makedirs(extraction_path, exist_ok=True) | |
with zipfile.ZipFile(zip_file_path, 'r') as zip_ref: | |
folder_name = os.path.basename(zip_file_path).split('.')[0] | |
zip_ref.extractall(extraction_path) | |
source_path = os.path.join(extraction_path, folder_name) | |
destination_path = os.path.join(extraction_path, folder_name) | |
if os.path.exists(destination_path): | |
print(f"Error: Destination path '{destination_path}' already exists") | |
else: | |
shutil.move(source_path, destination_path) | |
# Example usage: | |
zip_file_path = 'finetuned_bert_sentiment_harsh.zip' # Path to your ZIP file | |
extraction_path = 'bert_model_sentiment_v1' # Destination folder for extraction | |
unzip_and_save(zip_file_path, extraction_path) | |
# Load the fine-tuned model and tokenizer | |
model_path = "bert_model_sentiment_v1/finetuned_bert_sentiment_harsh" | |
tokenizer_path = "bert_model_sentiment_v1/finetuned_bert_sentiment_harsh" | |
def load_model(): | |
model = DistilBertForSequenceClassification.from_pretrained(model_path) | |
tokenizer = DistilBertTokenizerFast.from_pretrained(tokenizer_path) | |
return model, tokenizer | |
model, tokenizer = load_model() | |
def predict_sentiment(text): | |
device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
model.to(device) | |
tokenized = tokenizer(text, truncation=True, padding=True, return_tensors='pt').to(device) | |
outputs = model(**tokenized) | |
probs = F.softmax(outputs.logits, dim=-1) | |
preds = torch.argmax(outputs.logits, dim=-1).item() | |
probs_max = probs.max().detach().cpu().numpy() | |
prediction = "Positive" if preds == 1 else "Negative" | |
return prediction, probs_max * 100 | |
st.title("Sentiment Analysis App") | |
text = st.text_area("Enter your text:") | |
if st.button("Predict Sentiment"): | |
if text: | |
sentiment, confidence = predict_sentiment(text) | |
st.write(f"Sentiment: {sentiment}") | |
st.write(f"Confidence: {confidence:.2f}%") | |
else: | |
st.write("Please enter some text.") |