2-2/self-asi

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 = []