import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

class CatGPT:
    def __init__(self, model_name="OrionStarAI/Orion-14B"):
        self.model_name = model_name
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForCausalLM.from_pretrained(model_name)

    def generate(self, prompt, max_length=100):
        inputs = self.tokenizer(prompt, return_tensors="pt")
        outputs = self.model.generate(inputs.input_ids, max_length=max_length)
        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        return response

# Instantiate CatGPT
catgpt = CatGPT()

def catgpt_response(prompt):
    return catgpt.generate(prompt)

# Gradio Interface
iface = gr.Interface(
    fn=catgpt_response,
    inputs="text",
    outputs="text",
    title="CatGPT - Orion-14B Demo",
    description="This is a CatGPT demo using the Orion-14B model from OrionStarAI."
)

# Launch the app
iface.launch()