app.py

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


model_dir = "asadhu8/llama_3.2_1b_ddx_plus_medical"
tokenizer = AutoTokenizer.from_pretrained(model_dir)
model = AutoModelForCausalLM.from_pretrained(model_dir)

# Define a function for text generation
def generate_text(prompt, max_length=100, temperature=0.7):
    try:
        inputs = tokenizer.encode(prompt, return_tensors="pt")

        # Generate response using the model
        outputs = model.generate(
            inputs,
            max_length=max_length,
            temperature=temperature,
            pad_token_id=tokenizer.eos_token_id  # Handles padding
        )

        # Decode the model's output tokens into text
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)

        return response

    except Exception as e:
        return f"Error: {str(e)}"

# Define the Gradio interface
demo = gr.Interface(
    fn=generate_text,  # Function to call
    inputs=[
        gr.Textbox(label="Input Prompt", placeholder="Type your input here..."),
        gr.Slider(50, 300, value=100, step=10, label="Max Length"),
        gr.Slider(0.1, 1.0, value=0.7, step=0.1, label="Temperature"),
    ],
    outputs=gr.Textbox(label="Generated Response"),
    title="Local Hugging Face Model",
    description="Interact with a locally stored Hugging Face model for text generation.",
)

# Launch the Gradio app
demo.launch()