import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load the tokenizer and model from your Hugging Face repository
model_path = "sakthi54321/Power_chat_ai"  # Update with your model's repo

tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path)

# Function to generate a response based on input
def generate_response(prompt):
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(
        inputs["input_ids"], 
        max_length=150, 
        num_return_sequences=1, 
        do_sample=True, 
        temperature=0.7
    )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Function to determine if the prompt is code-related
def is_code_prompt(prompt):
    code_keywords = ['def', 'import', 'print', 'class', 'function', 'variable', 'while', 'for']
    return any(keyword in prompt for keyword in code_keywords)

# Combined response function (handles both general and code)
def handle_prompt(prompt):
    if is_code_prompt(prompt):
        return generate_response(f"### Instruction: Write a Python code to solve the following problem: {prompt}\n### Response:")
    else:
        return generate_response(f"### Instruction: {prompt}\n### Response:")

# Define Gradio interface with input and output fields
iface = gr.Interface(
    fn=handle_prompt, 
    inputs="text", 
    outputs="text", 
    live=True, 
    title="TinyLlama Assistant", 
    description="Interact with the TinyLlama-1B model for general and coding tasks. Enter a prompt to get a response, either in general text or Python code."
)

# Launch the app
iface.launch()