import subprocess
import sys
import gradio as gr
from transformers import pipeline

# Function to ensure required libraries are installed
def install_and_import(package):
    try:
        __import__(package)
    except ImportError:
        subprocess.check_call([sys.executable, "-m", "pip", "install", package])

# Install transformers and gradio if not installed
install_and_import("transformers")
install_and_import("gradio")

# Load DeepSeek-R1 model pipeline
def load_pipeline():
    global pipe
    pipe = pipeline(
        "text-generation",
        model="deepseek-ai/DeepSeek-R1",
        trust_remote_code=True
    )
    print("Model loaded successfully.")

load_pipeline()

# Function to handle user prompts and generate responses
def generate_response(prompt):
    try:
        # Generate response from the model
        response = pipe(prompt, max_length=100, num_return_sequences=1)
        return response[0]["generated_text"]
    except Exception as e:
        return f"Error: {str(e)}"

# Gradio Interface
interface = gr.Interface(
    fn=generate_response,
    inputs=gr.Textbox(lines=2, placeholder="Enter your prompt here..."),
    outputs=gr.Textbox(label="Generated Text"),
    title="DeepSeek-R1 Text Generator",
    description="Generate text using the DeepSeek-R1 model hosted on Hugging Face.",
)

# Launch the Gradio app
if __name__ == "__main__":
    interface.launch()