Update app.py

app.py

@@ -1,25 +1,32 @@
 import streamlit as st
-import openai
-import os
-
-# Set your OpenAI API key here or set it as an environment variable
-openai.api_key = os.getenv("OPENAI_API_KEY", "sk-None-RVvMry6BpMfG8KxYS6hTT3BlbkFJSOfj2bCwA1EmTRrAHo5y")
-
-# Function to classify lines using GPT-3
-def classify_lines_with_gpt3(text):
-    response = openai.Completion.create(
-        engine="text-davinci-003",
-        prompt=f"Classify the following Finnish contract specifications into categories: Urakka sisältää: Urakka ei sisältää: Tilaajan velvoitteet:Käytäntöjen tarkennukset:Hintojen tarkennukset: Muu:.\n\n{text}\n\n",
-        max_tokens=1024,
-        n=1,
-        stop=None,
-        temperature=0.5,
-    )
-    
-    classified_text = response.choices[0].text.strip()
-    return classified_text
-
-st.title("Finnish Contract Specifications Categorizer with GPT-3")
+from transformers import BertTokenizer, BertForSequenceClassification
+import torch
+import torch.nn.functional as F
+
+# Load the tokenizer and model
+model_name = "TurkuNLP/bert-base-finnish-cased-v1"
+tokenizer = BertTokenizer.from_pretrained(model_name)
+model = BertForSequenceClassification.from_pretrained(model_name, num_labels=6)  # Assuming 6 categories
+
+# Define categories
+categories = ["Urakka sisältää", "Urakka ei sisältää", "Tilaajan velvoitteet", "Käytäntöjen tarkennukset", "Hintojen tarkennukset", "Muu"]
+
+# Function to classify lines
+def classify_lines(text):
+    lines = text.split("\n")
+    categorized_lines = {category: [] for category in categories}
+
+    for line in lines:
+        if line.strip():  # Skip empty lines
+            inputs = tokenizer(line, return_tensors="pt", padding=True, truncation=True, max_length=512)
+            outputs = model(**inputs)
+            probs = F.softmax(outputs.logits, dim=1)
+            predicted_category = torch.argmax(probs, dim=1).item()
+            categorized_lines[categories[predicted_category]].append(line)
+
+    return categorized_lines
+
+st.title("Finnish Contract Specifications Categorizer with TurkuNLP BERT")
 
 st.write("Enter the contract specifications in Finnish:")
 
@@ -28,9 +35,13 @@ contract_text = st.text_area("Contract Specifications (Finnish):", height=300)
 
 if st.button("Classify"):
     if contract_text:
-        classified_text = classify_lines_with_gpt3(contract_text)
+        categories = classify_lines(contract_text)
 
         st.write("Classified Contract Specifications:")
-        st.write(classified_text)
+
+        for category, lines in categories.items():
+            st.write(f"### {category}")
+            for line in lines:
+                st.write(f"- {line}")
     else:
-        st.write("Please enter the contract specifications.")
+        st.write("Please enter the contract specifications.")