import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

class RPG_Game:
    def __init__(self):
        self.model_name = "nomic-ai/gpt4all-j"
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
        self.model = AutoModelForCausalLM.from_pretrained(self.model_name)
        self.inventory = {}
        
    def generate_response(self, prompt):
        input_tokens = self.tokenizer.encode(prompt, return_tensors="pt")
        response = self.model.generate(input_tokens)
        response_text = self.tokenizer.decode(response[:, input_tokens.shape[-1]:][0], skip_special_tokens=True)
        return response_text
        
    def process_game_input(self, game_input):
        game_input = game_input.strip().lower()

        # Check for specific commands to manage the inventory
        if game_input.startswith("pick up") or game_input.startswith("get"):
            item = game_input.split(" ", 1)[-1].strip()
            if item not in self.inventory:
                self.inventory[item] = 1
            else:
                self.inventory[item] += 1
            return f"You have picked up {item}."
        elif game_input == "inventory":
            return "Your inventory: " + ", ".join(f"{k}: {v}" for k, v in self.inventory.items())

        # Generate the GPT4ALL response
        else:
            response = self.generate_response(game_input)
            return response

game = RPG_Game()

def gradio_interface(prompt: str):
    return game.process_game_input(prompt)

inputs = gr.inputs.Textbox(lines=5, label="Your Input")
outputs = gr.outputs.Textbox(label="Game Response")
iface = gr.Interface(fn=gradio_interface, inputs=inputs, outputs=outputs)

if __name__ == "__main__":
    iface.launch()