import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load DeepSeek-Coder-7B Model
MODEL_NAME = "deepseek-ai/deepseek-coder-7b-instruct"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME, 
    torch_dtype=torch.float16, 
    device_map="auto"
)

# System Prompt to Guide the Model
SYSTEM_PROMPT = """
You are a highly skilled AI specialized in programming and mathematics. 
- For coding questions, provide clear explanations and format code inside triple backticks.
- For math problems, explain step-by-step solutions neatly.
- Keep responses professional, concise, and well-structured.
"""

def generate_response(prompt):
    formatted_prompt = f"{SYSTEM_PROMPT}\n\nUser: {prompt}\nAI:"  # Injecting system instructions
    inputs = tokenizer(formatted_prompt, return_tensors="pt").to("cuda")
    outputs = model.generate(inputs["input_ids"], max_length=700)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    return response

# Create Gradio Interface
interface = gr.Interface(
    fn=generate_response,
    inputs=gr.Textbox(lines=5, placeholder="Enter your math or coding question here..."),
    outputs="text",
    title="DeepSeek Coder & Math Pro",
    description="Ask anything about programming or mathematics!",
    theme="default",
)

interface.launch()