Create CatGPT.py

CatGPT.py

@@ -0,0 +1,32 @@
+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()