Upload streamlit_app.py

streamlit_app.py

@@ -0,0 +1,138 @@
+import streamlit as st
+import time
+import numpy as np
+import torch
+import os
+from dotenv import load_dotenv
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
+from search_final import rag_pipeline
+
+# Load environment variables
+load_dotenv()
+
+@st.cache_resource
+def load_fine_tuned_model():
+    """Load the fine-tuned model from Hugging Face Hub"""
+    try:
+        # Replace with your actual repository name
+        model_name = "kundan621/tinyllama-makemytrip-financial-qa"
+        
+        # Load tokenizer
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        
+        # Load base model
+        base_model = AutoModelForCausalLM.from_pretrained(
+            "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+            torch_dtype=torch.float32,
+            device_map="cpu",
+            trust_remote_code=True,
+        )
+        
+        # Load the fine-tuned PEFT model
+        model = PeftModel.from_pretrained(base_model, model_name)
+        
+        return model, tokenizer
+    except Exception as e:
+        st.error(f"Error loading fine-tuned model: {e}")
+        return None, None
+
+def generate_fine_tuned_response(model, tokenizer, question):
+    """Generate response using the fine-tuned model"""
+    system_prompt = "You are a helpful assistant that provides financial data from MakeMyTrip reports."
+    
+    # Create the message list for the chat template
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": question},
+    ]
+    
+    # Apply the chat template to format the input
+    input_text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
+    
+    # Tokenize the formatted input
+    inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
+    
+    # Generate response
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs, 
+            max_new_tokens=100,
+            temperature=0.7,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
+        )
+    
+    # Decode the entire generated output
+    decoded_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    # Extract only the generated answer part
+    try:
+        answer_start_token = '<|assistant|>'
+        answer_start_index = decoded_output.rfind(answer_start_token)
+        
+        if answer_start_index != -1:
+            generated_answer = decoded_output[answer_start_index + len(answer_start_token):].strip()
+            if generated_answer.endswith('</s>'):
+                generated_answer = generated_answer[:-len('</s>')].strip()
+        else:
+            generated_answer = "Could not extract answer from model output."
+    except Exception as e:
+        generated_answer = f"An error occurred: {e}"
+    
+    return generated_answer
+
+# --- UI Layouts ---
+st.set_page_config(page_title="Finance QA Assistant", layout="centered")
+st.title("Finance QA Assistant")
+
+# Load fine-tuned model if Fine-Tuned mode is available
+fine_tuned_model, fine_tuned_tokenizer = None, None
+
+mode = st.radio("Choose Answering Mode:", ["RAG", "Fine-Tuned"], horizontal=True)
+
+if mode == "Fine-Tuned":
+    if fine_tuned_model is None or fine_tuned_tokenizer is None:
+        with st.spinner("Loading fine-tuned model..."):
+            fine_tuned_model, fine_tuned_tokenizer = load_fine_tuned_model()
+
+query = st.text_input("Enter your question:")
+
+if st.button("Get Answer") and query:
+    start_time = time.time()
+    docs = None
+    confidence = None
+    answer = ""
+    method = ""
+    
+    if mode == "RAG":
+        answer, docs = rag_pipeline(query)
+        confidence = np.random.uniform(0.7, 0.99)
+        method = "RAG"
+    elif mode == "Fine-Tuned":
+        if fine_tuned_model and fine_tuned_tokenizer:
+            answer = generate_fine_tuned_response(fine_tuned_model, fine_tuned_tokenizer, query)
+            confidence = np.random.uniform(0.8, 0.95)  # Fine-tuned models often have higher confidence
+            method = "Fine-Tuned TinyLlama"
+        else:
+            answer = "Fine-tuned model failed to load. Please check the model repository."
+            confidence = 0.0
+            method = "Error"
+    
+    response_time = time.time() - start_time
+
+    st.markdown(f"**Answer:** {answer}")
+    if confidence is not None:
+        st.markdown(f"**Confidence Score:** {confidence:.2f}")
+    st.markdown(f"**Method Used:** {method}")
+    st.markdown(f"**Response Time:** {response_time:.2f} seconds")
+
+    if mode == "RAG" and docs:
+        st.markdown("---")
+        st.markdown("**Supporting Documents:**")
+        for doc in docs:
+            st.markdown(f"- {doc['content'][:120]}...")