|
|
import os |
|
|
import gradio as gr |
|
|
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM |
|
|
|
|
|
|
|
|
MODEL_ID = "PuruAI/Medini_Intelligence" |
|
|
FALLBACK_MODEL = "gpt2" |
|
|
HF_TOKEN = os.getenv("HF_TOKEN") |
|
|
|
|
|
def load_model(model_id): |
|
|
"""Load Medini if available, otherwise fallback to GPT-2.""" |
|
|
try: |
|
|
print(f"🔹 Loading model: {model_id}") |
|
|
tokenizer = AutoTokenizer.from_pretrained(model_id, token=HF_TOKEN) |
|
|
model = AutoModelForCausalLM.from_pretrained(model_id, token=HF_TOKEN) |
|
|
return pipeline("text-generation", model=model, tokenizer=tokenizer) |
|
|
except Exception as e: |
|
|
print(f"❌ Failed to load {model_id}: {e}") |
|
|
print("⏩ Falling back to GPT-2 (no token needed)") |
|
|
return pipeline("text-generation", model=FALLBACK_MODEL) |
|
|
|
|
|
|
|
|
generator = load_model(MODEL_ID) |
|
|
|
|
|
def generate_text(prompt): |
|
|
outputs = generator(prompt, max_length=200, num_return_sequences=1) |
|
|
return outputs[0]["generated_text"] |
|
|
|
|
|
|
|
|
iface = gr.Interface( |
|
|
fn=generate_text, |
|
|
inputs="text", |
|
|
outputs="text", |
|
|
title="Medini Intelligence", |
|
|
description="Custom AI Agent with fallback to GPT-2" |
|
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
|
iface.launch() |
|
|
|