app.py

import os
import multiprocessing
import concurrent.futures
from langchain_community.document_loaders import TextLoader, DirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from transformers import AutoModel, AutoTokenizer
import torch.nn.functional as F
import faiss
import torch
import numpy as np
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
from datetime import datetime
import json
import gradio as gr
import re
from threading import Thread

class DocumentRetrievalAndGeneration:
    def __init__(self, embedding_model_name, lm_model_id, data_folder):
        self.all_splits = self.load_documents(data_folder)
        
        # Get token from HF Spaces environment
        hf_token = os.getenv('HF_TOKEN')
        print(f"Token found: {hf_token is not None}")
        
        self.embedding_tokenizer = AutoTokenizer.from_pretrained(embedding_model_name, token=hf_token)
        self.embedding_model = AutoModel.from_pretrained(embedding_model_name, token=hf_token)
        self.gpu_index = self.create_faiss_index()
        self.tokenizer, self.model = self.initialize_llm(lm_model_id)

    def load_documents(self, folder_path):
        loader = DirectoryLoader(folder_path, loader_cls=TextLoader)
        documents = loader.load()
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=5000, chunk_overlap=250)
        all_splits = text_splitter.split_documents(documents)
        print('Length of documents:', len(documents))
        print("LEN of all_splits", len(all_splits))
        for i in range(min(3, len(all_splits))):
            print(all_splits[i].page_content[:200] + "...")
        return all_splits
        
    def encode_texts(self, texts):
        encoded_input = self.embedding_tokenizer(texts, padding=True, truncation=True, max_length=512, return_tensors='pt')
        with torch.no_grad():
            model_output = self.embedding_model(**encoded_input)
            embeddings = self.mean_pooling(model_output, encoded_input['attention_mask'])
            embeddings = F.normalize(embeddings, p=2, dim=1)
        return embeddings.cpu().numpy()

    def mean_pooling(self, model_output, attention_mask):
        token_embeddings = model_output[0]
        input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
        return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)

    def create_faiss_index(self):
        all_texts = [split.page_content for split in self.all_splits]
    
        batch_size = 512  # Reduced for Spaces
        all_embeddings = []
        
        for i in range(0, len(all_texts), batch_size):
            batch_texts = all_texts[i:i+batch_size]
            batch_embeddings = self.encode_texts(batch_texts)
            all_embeddings.append(batch_embeddings)
            print(f"Processed batch {i//batch_size + 1}/{(len(all_texts) + batch_size - 1)//batch_size}")
        
        embeddings = np.vstack(all_embeddings)
        index = faiss.IndexFlatL2(embeddings.shape[1])
        index.add(embeddings)
        
        # Try GPU first, fallback to CPU if fails
        try:
            if torch.cuda.is_available():
                gpu_resource = faiss.StandardGpuResources()
                gpu_index = faiss.index_cpu_to_gpu(gpu_resource, 0, index)
                print("🚀 Using GPU for FAISS")
                return gpu_index
            else:
                print("💻 Using CPU for FAISS")
                return index
        except Exception as e:
            print(f"GPU FAISS failed: {e}, using CPU")
            return index

    def initialize_llm(self, model_id):
        quantization_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_use_double_quant=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16
        )
        
        hf_token = os.getenv('HF_TOKEN')
        print(f"LLM Token found: {hf_token is not None}") 
        print(f"Token starts with: {hf_token[:10] if hf_token else 'None'}...")  
        
        tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token)
        
        # Handle pad_token for latest transformers
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
            
        model = AutoModelForCausalLM.from_pretrained(
            model_id,
            torch_dtype=torch.bfloat16,
            device_map="auto",
            quantization_config=quantization_config,
            token=hf_token
        )
        
        print(f"🦙 Model loaded on: {model.device}")
        return tokenizer, model

    def generate_response_with_timeout(self, input_ids, max_new_tokens=800):
        try:
            streamer = TextIteratorStreamer(self.tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
            generate_kwargs = dict(
                input_ids=input_ids,
                max_new_tokens=max_new_tokens,
                do_sample=True,
                top_p=1.0,
                top_k=20,
                temperature=0.8,
                repetition_penalty=1.2,
                pad_token_id=self.tokenizer.eos_token_id,
                eos_token_id=self.tokenizer.eos_token_id,
                streamer=streamer,
            )
            
            thread = Thread(target=self.model.generate, kwargs=generate_kwargs)
            thread.start()
            
            generated_text = ""
            for new_text in streamer:
                generated_text += new_text
            
            thread.join()
            return generated_text
        except Exception as e:
            print(f"Error in generation: {str(e)}")
            return "Text generation process encountered an error"

    def query_and_generate_response(self, query):
        if not query.strip():
            return "Please enter a valid query", ""
            
        try:
            similarityThreshold = 1.0
            query_embedding = self.encode_texts([query])[0]
            distances, indices = self.gpu_index.search(np.array([query_embedding]), k=3)
            print("Distance", distances, "indices", indices)
            
            content = ""
            for idx, distance in zip(indices[0], distances[0]):
                content += "-" * 50 + "\n"
                content += self.all_splits[idx].page_content + "\n"
                print(f"📄 Chunk {idx} (distance: {distance:.3f})")

            conversation = [
                {"role": "system", "content": "You are a knowledgeable assistant with access to a comprehensive database."},
                {"role": "user", "content": f"""
                I need you to answer my question and provide related information in a specific format.
                I have provided five relatable json files {content}, choose the most suitable chunks for answering the query.
                RETURN ONLY SOLUTION without additional comments, sign-offs, retrived chunks, refrence to any Ticket or extra phrases. Be direct and to the point.
                IF THERE IS NO ANSWER RELATABLE IN RETRIEVED CHUNKS, RETURN "NO SOLUTION AVAILABLE".
                DO NOT GIVE REFRENCE TO ANY CHUNKS OR TICKETS,BE ON POINT.
                
                Here's my question:
                Query: {query}
                Solution==>
                """}
            ]
            
            input_ids = self.tokenizer.apply_chat_template(conversation, return_tensors="pt").to(self.model.device)
            
            start_time = datetime.now()
            generated_response = self.generate_response_with_timeout(input_ids)
            elapsed_time = datetime.now() - start_time

            print("Generated response:", generated_response)
            print("Time elapsed:", elapsed_time)

            solution_text = generated_response.strip()
            if "Solution:" in solution_text:
                solution_text = solution_text.split("Solution:", 1)[1].strip()

            # Post-processing to remove "assistant" prefix
            solution_text = re.sub(r'^assistant\s*', '', solution_text, flags=re.IGNORECASE)
            solution_text = solution_text.strip()

            return solution_text, content[:1000] + "..." if len(content) > 1000 else content
            
        except Exception as e:
            print(f"Error in query processing: {e}")
            return f"Error processing query: {str(e)}", ""

    def qa_infer_gradio(self, query):
        response = self.query_and_generate_response(query)
        return response

# Initialize the system
print("Initializing TI E2E Forum Assistant...")

embedding_model_name = 'sentence-transformers/all-MiniLM-L6-v2'  # More compatible model
lm_model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
data_folder = 'sample_embedding_folder2'  # Make sure this folder exists

try:
    doc_retrieval_gen = DocumentRetrievalAndGeneration(embedding_model_name, lm_model_id, data_folder)
    print("System initialized successfully!")
    
    # Your exact same CSS and examples
    css_code = """
        .gradio-container {
            background-color: #daccdb;
        }
        button {
            background-color: #927fc7;
            color: black;
            border: 1px solid black;
            padding: 10px;
            margin-right: 10px;
            font-size: 16px;
            font-weight: bold;
        }
    """
    
    EXAMPLES = [
        "On which devices can the VIP and CSI2 modules operate simultaneously?", 
        "I'm using Code Composer Studio 5.4.0.00091 and enabled FPv4SPD16 floating point support for CortexM4 in TDA2. However, after building the project, the .asm file shows --float_support=vfplib instead of FPv4SPD16. Why is this happening?", 
        "Could you clarify the maximum number of cameras that can be connected simultaneously to the video input ports on the TDA2x SoC, considering it supports up to 10 multiplexed input ports and includes 3 dedicated video input modules?"
    ]

    interface = gr.Interface(
        fn=doc_retrieval_gen.qa_infer_gradio,
        inputs=[gr.Textbox(label="QUERY", placeholder="Enter your query here", lines=3)],
        allow_flagging='never',
        examples=EXAMPLES,
        cache_examples=False,
        outputs=[
            gr.Textbox(label="RESPONSE", lines=8),
            gr.Textbox(label="RELATED QUERIES", lines=5)
        ],
        css=css_code,
        title="🤖 TI E2E FORUM",
        description="Ask technical questions and get answers based on the TI E2E knowledge base"
    )

    # Launch with public link for Spaces
    interface.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )
    
except Exception as e:
    print(f"Failed to initialize: {e}")
    
    # Fallback simple interface
    def fallback_response(query):
        return "System initialization failed. Please check the logs.", ""
    
    fallback_interface = gr.Interface(
        fn=fallback_response,
        inputs=[gr.Textbox(label="QUERY")],
        outputs=[gr.Textbox(label="ERROR"), gr.Textbox(label="INFO")],
        title="TI E2E FORUM - Initialization Error"
    )
    
    fallback_interface.launch(server_name="0.0.0.0", server_port=7860)