chanderbawa1983/factual-feed

1"""2Hugging Face Transformers Processor - Free local LLM processing3Supports models like Llama, Mistral, FLAN-T5, etc.4"""5 6import torch7from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline8import json9import logging10from typing import List, Dict, Optional11from prompt_templates import PromptTemplates12 13logging.basicConfig(level=logging.INFO)14logger = logging.getLogger(__name__)15 16class HuggingFaceProcessor:17    def __init__(self, model_name: str = "microsoft/DialoGPT-medium"):18        """19        Initialize Hugging Face processor20        21        Args:22            model_name: Model to use (e.g., "microsoft/DialoGPT-medium", "google/flan-t5-base")23        """24        self.model_name = model_name25        self.prompt_templates = PromptTemplates()26        self.device = "cuda" if torch.cuda.is_available() else "cpu"27        28        # Initialize model and tokenizer29        self.tokenizer = None30        self.model = None31        self.pipeline = None32        33        self._load_model()34    35    def _load_model(self):36        """Load the model and tokenizer"""37        try:38            logger.info(f"Loading model {self.model_name} on {self.device}...")39            40            # For text generation models41            if "flan-t5" in self.model_name.lower():42                self.pipeline = pipeline(43                    "text2text-generation",44                    model=self.model_name,45                    device=0 if self.device == "cuda" else -1,46                    max_length=51247                )48            else:49                self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)50                self.model = AutoModelForCausalLM.from_pretrained(51                    self.model_name,52                    torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,53                    device_map="auto" if self.device == "cuda" else None54                )55                56                # Add padding token if not present57                if self.tokenizer.pad_token is None:58                    self.tokenizer.pad_token = self.tokenizer.eos_token59            60            logger.info(f"Model loaded successfully on {self.device}")61            62        except Exception as e:63            logger.error(f"Error loading model: {e}")64            logger.info("Falling back to smaller model...")65            try:66                self.model_name = "microsoft/DialoGPT-small"67                self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)68                self.model = AutoModelForCausalLM.from_pretrained(self.model_name)69                if self.tokenizer.pad_token is None:70                    self.tokenizer.pad_token = self.tokenizer.eos_token71                logger.info("Fallback model loaded successfully")72            except Exception as e2:73                logger.error(f"Failed to load fallback model: {e2}")74    75    def _generate_text(self, prompt: str, max_length: int = 200) -> str:76        """Generate text using the loaded model"""77        try:78            if self.pipeline:79                # Use pipeline for T5-style models80                result = self.pipeline(prompt, max_length=max_length, do_sample=True, temperature=0.7)81                return result[0]['generated_text'].strip()82            83            elif self.model and self.tokenizer:84                # Use model directly for GPT-style models85                inputs = self.tokenizer(86                    prompt,87                    return_tensors="pt",88                    truncation=True,89                    max_length=51290                )91                input_ids = inputs["input_ids"]92                if self.device == "cuda":93                    input_ids = input_ids.to(self.device)94                95                with torch.no_grad():96                    outputs = self.model.generate(97                        input_ids,98                        max_length=input_ids.shape[1] + max_length,99                        num_return_sequences=1,100                        temperature=0.7,101                        do_sample=True,102                        pad_token_id=self.tokenizer.eos_token_id103                    )104                105                generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)106                # Remove the original prompt from the response if present107                if generated_text.startswith(prompt):108                    generated_text = generated_text[len(prompt):]109                return generated_text.strip()110            111            else:112                logger.error("No model available for generation")113                return ""114                115        except Exception as e:116            logger.error(f"Error generating text: {e}")117            return ""118    119    def generate_summary(self, article: Dict) -> str:120        """Generate article summary"""121        try:122            # Simplified prompt for better results with smaller models123            prompt = f"Summarize this news article in 3-4 sentences:\n\nTitle: {article.get('title', '')}\n\nArticle: {article['content'][:1000]}...\n\nSummary:"124            125            summary = self._generate_text(prompt, max_length=150)126            127            if summary and len(summary) > 20:128                logger.info(f"Generated summary for article: {article.get('title', 'Untitled')[:50]}...")129                return summary130            else:131                # Fallback to extractive summary132                sentences = article['content'].split('. ')133                return '. '.join(sentences[:3]) + '.'134                135        except Exception as e:136            logger.error(f"Error generating summary: {e}")137            return f"Summary unavailable. Original content: {article['content'][:200]}..."138    139    def generate_draft_trivia(self, article: Dict, entities: List[Dict]) -> List[Dict]:140        """Generate draft trivia questions"""141        try:142            # Simplified approach for smaller models143            questions = []144            145            # Take top 3 entities for trivia146            top_entities = entities[:3]147            148            for entity in top_entities:149                entity_text = entity['text']150                entity_type = entity['label']151                context = entity.get('context_sentence', '')152                153                # Simple question generation154                if entity_type == 'PERSON':155                    question_text = f"Who is mentioned in the article in relation to {context[:50]}...?"156                elif entity_type == 'ORG':157                    question_text = f"Which organization is mentioned in the article?"158                elif entity_type == 'GPE':159                    question_text = f"Which location is mentioned in the article?"160                elif entity_type in ['DATE', 'CARDINAL', 'MONEY', 'PERCENT']:161                    question_text = f"What {entity_type.lower()} is mentioned in the article?"162                else:163                    question_text = f"What {entity_type.lower()} is mentioned in the article?"164                165                # Generate simple distractors based on entity type166                if entity_type == 'PERSON':167                    distractors = ["John Smith", "Jane Doe", "Michael Johnson"]168                elif entity_type == 'ORG':169                    distractors = ["Microsoft", "Google", "Amazon"]170                elif entity_type == 'GPE':171                    distractors = ["New York", "London", "Tokyo"]172                elif entity_type == 'DATE':173                    distractors = ["2023", "2022", "2021"]174                elif entity_type == 'MONEY':175                    distractors = ["$1 million", "$5 million", "$10 million"]176                else:177                    distractors = ["Option A", "Option B", "Option C"]178                179                question = {180                    'question': question_text,181                    'correct_answer': entity_text,182                    'source_sentence': context or f"The article mentions {entity_text}.",183                    'distractors': distractors,184                    'difficulty': 'medium',185                    'entity_type': entity_type,186                    'article_url': article.get('url', ''),187                    'article_title': article.get('title', ''),188                    'source_article': article.get('content', '')189                }190                191                questions.append(question)192            193            logger.info(f"Generated {len(questions)} draft trivia questions")194            return questions195            196        except Exception as e:197            logger.error(f"Error generating trivia: {e}")198            return []199    200    def verify_fact(self, claim: str, source_text: str) -> Dict:201        """Simple fact verification"""202        try:203            # Simple keyword-based verification for smaller models204            claim_lower = claim.lower()205            source_lower = source_text.lower()206            207            # Extract key terms from claim208            claim_words = set(claim_lower.split())209            source_words = set(source_lower.split())210            211            # Calculate overlap212            overlap = len(claim_words.intersection(source_words))213            total_claim_words = len(claim_words)214            215            if total_claim_words == 0:216                confidence = 0217            else:218                confidence = overlap / total_claim_words219            220            if confidence > 0.7:221                status = "SUPPORTED"222                explanation = f"High word overlap ({confidence:.1%}) between claim and source"223            elif confidence > 0.4:224                status = "PARTIALLY_SUPPORTED"225                explanation = f"Moderate word overlap ({confidence:.1%}) between claim and source"226            else:227                status = "NOT_SUPPORTED"228                explanation = f"Low word overlap ({confidence:.1%}) between claim and source"229            230            return {231                'status': status,232                'explanation': explanation,233                'relevant_quote': source_text[:200] + "..." if len(source_text) > 200 else source_text234            }235            236        except Exception as e:237            logger.error(f"Error verifying fact: {e}")238            return {239                'status': 'ERROR',240                'explanation': f'Verification failed: {str(e)}',241                'relevant_quote': ''242            }243    244    @staticmethod245    def get_recommended_models() -> Dict[str, str]:246        """Get recommended models for different use cases"""247        return {248            "Small/Fast": "microsoft/DialoGPT-small",249            "Medium": "microsoft/DialoGPT-medium", 250            "Text2Text": "google/flan-t5-base",251            "Large (GPU)": "microsoft/DialoGPT-large",252            "Instruction": "google/flan-t5-large"253        }254