region:us
import sys
import subprocess
import platform
import warnings
warnings.filterwarnings("ignore")

import torch
from torch import nn
from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
import gym

class SimpleNavigationEnv(gym.Env):
    """Simple Navigation Environment"""

    metadata = {'render.modes': ['ansi']}

    def __init__(self):
        super().__init__()

        # Initialize environment variables
        self.env = gym.make('CartPole-v1')
        self.action_space = self.env.action_space
        self.observation_space = self.env.observation_space
        
        # Load tokenizer and LLM
        tokenizer = AutoTokenizer.from_pretrained("xlm-roberta-base")
        llm =pipeline('text-classification', model="xlm-roberta-base", tokenizer="xlm-roberta-base")
    
        # Define helper methods for interacting with the environment    
        def reset():
            observation = self.env.reset()
            return observation
    
        def render(mode='ansi'):
            return self.env.render(mode)
    
        def step(action):
            obs, reward, done, info = self.env.step(action)
            return obs, reward, done, info
    
        # Define the agent's behaviour    
        def interact(question):
            # Tokenize question
            tokens = tokenizer([question], padding=True, truncation=True, return_tensors='pt')

            # Classify question using LLM
            classifications = llm(tokens['input_ids'])
            
            # Based on the classifier's output, decide on the next action    
            if classifications[0]['label']=='go left':
                action = self.env.actions[0]
                
            elif classifications[0]['label']=='go right':
                action = self.env.actions[1]
                
            else:
                raise ValueError(f"Invalid command '{classifications[0]['label']}'")
        
            # Perform the selected action and observe the new state    
            obs, reward, done, info = step(action)
            return obs, reward, done, info
    
    def close(self):
        self.env.close()
        
class HybridAgent:
    def __init__(self):
        self.model_path = "hybrid_model.bin"
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        self.num_labels = 2
        self.batch_size = 16
        self.lr = 2e-5
        self.epochs = 10
        self.max_seq_len = 128
        self.save_freq = 1
        
        # Check if necessary dependencies are installed, otherwise install them
        deps = ["transformers", "sentencepiece", "numpy"]
        for dep in deps:
            if dep not in sys.modules:
                print(f"{dep} not found. Installing...")
                pip_cmd = ["pip", "install", dep]
                subprocess.run(pip_cmd)
                    
        # Initialize the LLM and TinyLM MT models
        self.tokenizer = AutoTokenizer.from_pretrained("xlm-roberta-base")
        self.llm = pipeline('text-classification', model="xlm-roberta-base", tokenizer="xlm-roberta-base")
        self.tinylmmt = AutoModelForSequenceClassification.from_pretrained("cardiffnlp/twitter-xlm-roberta-large-sentiment").to(self.device)
        
        # Load the training data and split it into train/validation sets
        df = pd.read_csv('./data/conversation_logs.csv')
        self.train_df, self.val_df = train_test_split(df, test_size=.2)
        
    def fit(self):
        # Fine tune the combined model on the training data
        optimizer = torch.optim.AdamW(list(self.tinylmmt.parameters()) + list(self.llm.model.classifier.parameters()), lr=self.lr)
        loss_fn = nn.CrossEntropyLoss()
        
        for epoch in range(self.epochs):
            self.train_epoch(optimizer, loss_fn)
            self.evaluate()
            if (epoch+1)%self.save_freq==0:
                self.save()
                
    def train_epoch(self, optimizer, loss_fn):
        self.tinylmmt.train()
        self.llm.model.eval()
        
        losses = []