Spaces:
Running
Running
| from langchain_community.embeddings.fastembed import FastEmbedEmbeddings | |
| from langchain_groq import ChatGroq | |
| import os | |
| import json | |
| from typing import List, Dict | |
| class LLMProcessor: | |
| def __init__(self): | |
| """Initialize embedding model and Groq LLM""" | |
| self.api_key = os.getenv("GROQ_API_KEY") | |
| # Use FastEmbed instead of SentenceTransformer | |
| self.embed_model = FastEmbedEmbeddings() | |
| self.llm = ChatGroq( | |
| model_name="mixtral-8x7b-32768", | |
| api_key=self.api_key | |
| ) | |
| def format_context(self, chunks: List[Dict]) -> str: | |
| """Format retrieved chunks into a structured context for the LLM""" | |
| context_parts = [] | |
| for chunk in chunks: | |
| try: | |
| headings = json.loads(chunk['headings']) | |
| if headings: | |
| context_parts.append(f"Section: {' > '.join(headings)}") | |
| except: | |
| pass | |
| if chunk['page']: | |
| context_parts.append(f"Page {chunk['page']}:") | |
| context_parts.append(chunk['text']) | |
| context_parts.append("-" * 40) | |
| return "\n".join(context_parts) | |
| def generate_answer(self, context: str, question: str) -> str: | |
| """Generate answer using structured context""" | |
| prompt = f"""Based on the following excerpts from a document: | |
| {context} | |
| Please answer this question: {question} | |
| Make use of the section information and page numbers in your answer when relevant. | |
| """ | |
| return self.llm.invoke(prompt) | |