Update app.py

app.py

@@ -1,23 +1,12 @@
 import os
 import pandas as pd
+import streamlit as st
 import transformers
 import torch
-from transformers import Trainer, TrainingArguments, DataCollatorForLanguageModeling, AutoTokenizer, AutoModelForCausalLM
+from transformers import Trainer, TrainingArguments, DataCollatorForLanguageModeling, AutoTokenizer, AutoModelForCausalLM, pipeline
 from torch.utils.data import Dataset
-import streamlit as st
-
-# Path to the directory where Ollama stores models
-model_dir = "C:/Users/myuser/.ollama/models/meta-llama-3-8b"
 
-# Load the tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained(model_dir)
-model = AutoModelForCausalLM.from_pretrained(model_dir)
-
-# Add a padding token to the tokenizer if it doesn't have one
-if tokenizer.pad_token is None:
-    tokenizer.add_special_tokens({'pad_token': '[PAD]'})
-
-# Custom Dataset class
+# Define the Custom Dataset class
 class CustomTextDataset(Dataset):
     def __init__(self, tokenized_inputs):
         self.input_ids = tokenized_inputs['input_ids']
@@ -38,64 +27,77 @@ def prepare_dataset(texts, tokenizer, block_size=128):
     dataset = CustomTextDataset(inputs)
     return dataset
 
+# Function to fine-tune the model
+def fine_tune_model(train_dataset, model, tokenizer):
+    data_collator = DataCollatorForLanguageModeling(
+        tokenizer=tokenizer,
+        mlm=False,
+    )
+
+    training_args = TrainingArguments(
+        output_dir="./results",
+        overwrite_output_dir=True,
+        num_train_epochs=3,
+        per_device_train_batch_size=4,
+        save_steps=10_000,
+        save_total_limit=2,
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        data_collator=data_collator,
+        train_dataset=train_dataset,
+    )
+
+    trainer.train()
+
 # Load the dataset
-file_path = "path/to/it_support_transcript_dataset.csv"
-df = pd.read_csv(file_path)
-
-# Filter the dataset based on the given criteria and make a copy
-filtered_df = df[
-    (df['Resolution Status'] == 'Resolved') &
-    (df['Customer Satisfaction (CSAT) Score'] >= 4) &
-    (df['Customer Feedback Comments'].isin(['Very satisfied', 'Satisfied']))
-].copy()
-
-# Combine only the interaction notes into a single text for training
-filtered_df.loc[:, 'training_text'] = filtered_df['Interaction Notes']
-
-# Select the training text
-training_texts = filtered_df['training_text'].tolist()
-
-# Create CustomTextDataset for fine-tuning
-train_dataset = prepare_dataset(training_texts, tokenizer)
-
-# Data collator for language modeling
-data_collator = DataCollatorForLanguageModeling(
-    tokenizer=tokenizer,
-    mlm=False,
-)
-
-# Training arguments
-training_args = TrainingArguments(
-    output_dir="./results",
-    overwrite_output_dir=True,
-    num_train_epochs=3,
-    per_device_train_batch_size=4,
-    save_steps=10_000,
-    save_total_limit=2,
-)
-
-# Trainer
-trainer = Trainer(
-    model=model,
-    args=training_args,
-    data_collator=data_collator,
-    train_dataset=train_dataset,
-)
-
-# Fine-tune the model
-trainer.train()
-
-# Streamlit app
-st.title("IT Support Assistant")
-
-# Create a text generation pipeline
-text_gen_pipeline = transformers.pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer
-)
+@st.cache
+def load_data(it_support_transcript_dataset.csv):
+    df = pd.read_csv(it_support_transcript_dataset.csv)
+    filtered_df = df[
+        (df['Resolution Status'] == 'Resolved') &
+        (df['Customer Satisfaction (CSAT) Score'] >= 4) &
+        (df['Customer Feedback Comments'].isin(['Very satisfied', 'Satisfied']))
+    ].copy()
+    filtered_df.loc[:, 'training_text'] = filtered_df['Interaction Notes']
+    training_texts = filtered_df['training_text'].tolist()
+    return training_texts
+
+# Streamlit UI
+st.title("IT Support Assistant - Training and Deployment")
+
+# File upload
+file_path = st.text_input("it_support_transcript_dataset.csv")
+train_button = st.button("Train Model")
+
+if file_path and train_button:
+    with st.spinner("Loading data and training the model..."):
+        training_texts = load_data(file_path)
+        
+        # Load the tokenizer and model from Hugging Face
+        model_name = "meta-llama/Meta-Llama-2-7B-chat-hf" # Use the available Llama model
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForCausalLM.from_pretrained(model_name)
+        
+        if tokenizer.pad_token is None:
+            tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+
+        train_dataset = prepare_dataset(training_texts, tokenizer)
+        fine_tune_model(train_dataset, model, tokenizer)
+
+    st.success("Model trained successfully!")
+
+# Interactive interface
+st.title("IT Support Assistant - Interaction")
 
 def generate_response(input_text):
+    text_gen_pipeline = pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer
+    )
     outputs = text_gen_pipeline(input_text, max_length=150, num_return_sequences=1)
     response = outputs[0]['generated_text']
     return response