raz-1412's picture
Update app.py
d9ac89f verified
Raw
History Blame Contribute Delete
2.41 kB
import json
from transformers import pipeline
import docx
from bs4 import BeautifulSoup
import streamlit as st
# Load the Hugging Face NER pipeline using a pre-trained model
nlp = pipeline("ner", model="dbmdz/bert-large-cased-finetuned-conll03-english")
def extract_text_from_file(uploaded_file):
"""Extract text from various file formats."""
file_type = uploaded_file.name.split(".")[-1].lower()
if file_type == "txt":
text = uploaded_file.getvalue().decode("utf-8")
elif file_type == "html":
soup = BeautifulSoup(uploaded_file.getvalue(), "html.parser")
text = soup.get_text()
elif file_type == "json":
data = json.load(uploaded_file)
text = "\n".join(str(value) for value in data.values()) if isinstance(data, dict) else json.dumps(data)
elif file_type == "docx":
doc = docx.Document(uploaded_file)
text = "\n".join([para.text for para in doc.paragraphs])
else:
raise ValueError("Unsupported file format")
return text
def infer_relationships(chat_data):
"""Process large text inputs in chunks to avoid memory issues."""
relationship_scores = []
max_chars = 1000 # Keep within the model's safe limits
# Split the text into chunks if it's too long
for i in range(0, len(chat_data), max_chars):
chunk = chat_data[i:i + max_chars]
relationship_scores.append(analyze_relationships(chunk))
return relationship_scores
def analyze_relationships(text_chunk):
"""Analyze relationships (NER) in text using the Hugging Face model."""
# Run Named Entity Recognition
entities = nlp(text_chunk)
return {"entities": [(ent['word'], ent['entity']) for ent in entities]}
def main():
st.title("Chat Relationship Analyzer")
uploaded_file = st.file_uploader("Upload a file", type=["txt", "html", "json", "docx"])
if uploaded_file is not None:
try:
chat_text = extract_text_from_file(uploaded_file)
# If the text is too long, warn the user and process in chunks
if len(chat_text) > 1000:
st.warning("Text is too long. Processing in chunks.")
relationship_scores = infer_relationships(chat_text)
st.json(relationship_scores)
except Exception as e:
st.error(f"Error: {e}")
if __name__ == "__main__":
main()