Update app.py

app.py

@@ -1,154 +1,163 @@
-# Configure page settings (MUST BE FIRST STREAMLIT COMMAND)
-import streamlit as st
-from streamlit_option_menu import option_menu
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
-from PyPDF2 import PdfReader
-
-# Set page config
-st.set_page_config(
-    page_title="Disease Analysis GPT",
-    layout="wide",
-    initial_sidebar_state="expanded"
-)
-
-# Load Hugging Face models and tokenizer for text generation
-@st.cache_resource
-def load_model():
-    tokenizer = AutoTokenizer.from_pretrained("harishussain12/Disease_Managment")
-    model = AutoModelForCausalLM.from_pretrained("harishussain12/Disease_Managment")
-    return tokenizer, model
-
-# Function to create a text generation pipeline
-@st.cache_resource
-def create_pipeline():
-    tokenizer, model = load_model()
-    return pipeline("text-generation", model=model, tokenizer=tokenizer)
-
-# Function to extract text from PDF file
-def read_pdf(file):
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+from accelerate import Accelerator
+
+# Check if GPU is available for better performance
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
+
+# Initialize the Accelerator for optimized inference
+accelerator = Accelerator()
+
+# Load models and tokenizers with FP16 for speed optimization if GPU is available
+model_dirs = [
+    "Poonawala/gpt2",
+    "Poonawala/MiriFur",
+    "Poonawala/Llama-3.2-1B"
+]
+
+models = {}
+tokenizers = {}
+
+def load_model(model_dir):
+    model = AutoModelForCausalLM.from_pretrained(model_dir, torch_dtype=torch.float16 if device.type == "cuda" else torch.float32)
+    tokenizer = AutoTokenizer.from_pretrained(model_dir)
+
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    # Move model to GPU/CPU as per availability
+    model = model.to(device)
+    return model, tokenizer
+
+# Load all models
+for model_dir in model_dirs:
+    model_name = model_dir.split("/")[-1]
     try:
-        reader = PdfReader(file)
-        text = ""
-        for page in reader.pages:
-            text += page.extract_text()
-        return text
+        model, tokenizer = load_model(model_dir)
+        models[model_name] = model
+        tokenizers[model_name] = tokenizer
+
+        # Batch warm-up inference to reduce initial response time
+        dummy_inputs = ["Hello", "What is a recipe?", "Explain cooking basics"]
+        for dummy_input in dummy_inputs:
+            input_ids = tokenizer.encode(dummy_input, return_tensors='pt').to(device)
+            with torch.no_grad():
+                model.generate(input_ids, max_new_tokens=1)
+
+        print(f"Loaded model and tokenizer from {model_dir}.")
     except Exception as e:
-        return f"Error reading PDF: {e}"
-
-# Load pipelines
-text_pipeline = create_pipeline()
-
-# Custom CSS for styling
-st.markdown(
-    """
-    <style>
-    body {
-        font-family: 'Arial', sans-serif;
-    }
-    .stButton button {
-        background-color: #0b2545;
-        color: white;
-        border: none;
-        border-radius: 25px;
-        padding: 8px 20px;
-        font-size: 14px;
-        font-weight: bold;
-        cursor: pointer;
-    }
-    .stButton button:hover {
-        background-color: #0a1b35;
+        print(f"Failed to load model from {model_dir}: {e}")
+        continue
+
+def get_response(prompt, model_name, user_type):
+    if model_name not in models:
+        return "Model not loaded correctly."
+
+    model = models[model_name]
+    tokenizer = tokenizers[model_name]
+
+    # Define different prompt templates based on user type
+    user_type_templates = {
+        "Expert": f"As an Expert, {prompt}\nAnswer:",
+        "Intermediate": f"As an Intermediate, {prompt}\nAnswer:",
+        "Beginner": f"Explain in simple terms: {prompt}\nAnswer:",
+        "Professional": f"As a Professional, {prompt}\nAnswer:"
     }
-    .search-box {
-        border-radius: 20px;
-        border: 1px solid #ccc;
-        padding: 10px;
-        width: 100%;
-        font-size: 16px;
-        background-color: #ffffff;
-    }
-    .info-box {
-        background-color: #f8f9fa;
-        border-left: 5px solid #0b2545;
-        padding: 15px;
-        border-radius: 5px;
-        font-size: 14px;
-    }
-    </style>
-    """,
-    unsafe_allow_html=True
-)
-
-# Sidebar
-with st.sidebar:
-    new_chat_button = st.button("New Chat", key="new_chat", help="Start a new chat to ask a different question.")
-    if new_chat_button:
-        st.session_state.clear()  # Clear session state to simulate a new chat
-
-    selected = option_menu(
-        menu_title=None,
-        options=[" Home", " Discover"],
-        icons=["house", "search"],
-        menu_icon="cast",
-        default_index=0,
-        styles={
-            "container": {"padding": "0!important", "background-color": "#3e4a5b"},
-            "icon": {"color": "#ffffff", "font-size": "16px"},
-            "nav-link": {
-                "font-size": "15px", 
-                "text-align": "left", 
-                "margin": "0px", 
-                "color": "#ffffff",
-                "font-weight": "bold",
-                "padding": "10px 20px",
-            },
-            "nav-link-selected": {"background-color": "#0b2545", "color": "white"},
-        }
-    )
-
-# Main content
-col1, col2, col3 = st.columns([1, 2, 1])
-
-with col2:
-    st.markdown("<h1 style='text-align: center;'>Disease Analysis GPT</h1>", unsafe_allow_html=True)
-    st.markdown("<h3 style='text-align: center;'>What do you want to know?</h3>", unsafe_allow_html=True)
-
-    # Model selection (now including Document Analysis)
-    model_selection = st.selectbox(
-        "Select a model", 
-        options=["Disease Analysis", "Document Analysis"], 
-        index=0
-    )
-
-    # If the user selects Document Analysis, show an error and prompt them to switch to Disease Analysis
-    if model_selection == "Document Analysis":
-        st.error("Please switch to 'Disease Analysis' model for generating responses. Document Analysis is not available in this version.")
-
-    # Search box
-    search_input = st.text_input(
-        "", 
-        placeholder="Type your question here...", 
-        label_visibility="collapsed",
-        help="Ask anything related to disease management."
-    )
-
-    # File upload below search box
-    uploaded_file = st.file_uploader("Upload a PDF file", type="pdf", help="Attach relevant files or documents to your query.")
-
-    if search_input:
-        with st.spinner("Generating response..."):
-            try:
-                if model_selection == "Disease Analysis":
-                    context = ""
-                    if uploaded_file is not None:
-                        file_content = read_pdf(uploaded_file)
-                        if "Error" in file_content:
-                            st.error(file_content)
-                        else:
-                            context = file_content
-
-                    query_input = search_input + (f"\n\nContext:\n{context}" if context else "")
-                    response = text_pipeline(query_input, max_length=200, num_return_sequences=1)
-                    st.markdown(f"### Response:\n{response[0]['generated_text']}")
-
-            except Exception as e:
-                st.error(f"Error generating response: {str(e)}")
+
+    # Get the appropriate prompt based on user type
+    prompt_template = user_type_templates.get(user_type, f"{prompt}\nAnswer:")
+
+    encoding = tokenizer(
+        prompt_template,
+        return_tensors='pt',
+        padding=True,
+        truncation=True,
+        max_length=500  # Increased length for larger inputs
+    ).to(device)
+
+    max_new_tokens = 200  # Increased to allow full-length answers
+
+    with torch.no_grad():
+        output = model.generate(
+            input_ids=encoding['input_ids'],
+            attention_mask=encoding['attention_mask'],
+            max_new_tokens=max_new_tokens,  # Higher value for longer answers
+            num_beams=3,         # Using beam search for better quality answers
+            repetition_penalty=1.2,  # Increased to reduce repetitive text
+            temperature=0.9,     # Slightly higher for creative outputs
+            top_p=0.9,           # Including more tokens for diverse generation
+            early_stopping=True,
+            pad_token_id=tokenizer.pad_token_id
+        )
+
+    response = tokenizer.decode(output[0], skip_special_tokens=True)
+    return response.strip()
+
+def process_input(prompt, model_name, user_type):
+    if prompt and prompt.strip():
+        return get_response(prompt, model_name, user_type)
+    else:
+        return "Please provide a prompt."
+
+# Gradio Interface with Modern Design
+with gr.Blocks(css="""
+body {
+    background-color: #faf3e0; /* Beige for a warm food-related theme */
+    font-family: 'Arial, sans-serif';
+}
+.title {
+    font-size: 2.5rem;
+    font-weight: bold;
+    color: #ff7f50; /* Coral color for a food-inspired look */
+    text-align: center;
+    margin-bottom: 1rem;
+}
+.container {
+    max-width: 900px;
+    margin: auto;
+    padding: 2rem;
+    background-color: #ffffff;
+    border-radius: 10px;
+    box-shadow: 0 4px 15px rgba(0, 0, 0, 0.1);
+}
+.button {
+    background-color: #ff7f50; /* Coral color for buttons */
+    color: white;
+    padding: 0.8rem 1.5rem;
+    font-size: 1rem;
+    border: none;
+    border-radius: 5px;
+    cursor: pointer;
+}
+.button:hover {
+    background-color: #ffa07a; /* Light salmon for hover effect */
+}
+""") as demo:
+
+    gr.Markdown("<div class='title'>Cookspert: Your Cooking Assistant</div>")
+
+    user_types = ["Expert", "Intermediate", "Beginner", "Professional"]
+
+    with gr.Tabs():
+        with gr.TabItem("Ask a Cooking Question"):
+            with gr.Row():
+                with gr.Column(scale=2):
+                    prompt = gr.Textbox(label="Ask about any recipe", placeholder="Ask question related to cooking here...", lines=2)
+                    model_name = gr.Radio(label="Choose Model", choices=list(models.keys()), interactive=True)
+                    user_type = gr.Dropdown(label="User Type", choices=user_types, value="Beginner")
+                    submit_button = gr.Button("ChefGPT", elem_classes="button")
+
+            response = gr.Textbox(
+                label="🍽️ Response",
+                placeholder="Your answer will appear here...",
+                lines=10,
+                interactive=False,
+                show_copy_button=True
+            )
+
+        submit_button.click(fn=process_input, inputs=[prompt, model_name, user_type], outputs=response)
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", share=True, debug=True)