Create model.py

model.py

@@ -0,0 +1,214 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import Dataset, DataLoader
+from datasets import load_dataset
+from transformers import AutoTokenizer
+from tqdm import tqdm
+import math
+import speech_recognition as sr
+import pyttsx3
+from googlesearch import search
+import warnings
+from typing import List, Dict, Union
+
+# Ignore warnings
+warnings.filterwarnings("ignore")
+
+class WebSearchWrapper:
+    """Wrapper for web search with caching"""
+    def __init__(self, cache_size: int = 100):
+        self.cache: Dict[str, List[str]] = {}
+        self.cache_size = cache_size
+        
+    def search(self, query: str, num_results: int = 3) -> List[str]:
+        """Perform web search with caching"""
+        if query.lower() in self.cache:
+            return self.cache[query.lower()]
+        
+        try:
+            search_results = list(search(query, num_results=num_results, stop=num_results, pause=2))
+            self._add_to_cache(query, search_results)
+            return search_results
+        except Exception as e:
+            print(f"Web search error: {e}")
+            return []
+    
+    def _add_to_cache(self, query: str, results: List[str]):
+        """Add results to cache with LRU eviction policy"""
+        if len(self.cache) >= self.cache_size:
+            self.cache.pop(next(iter(self.cache)))
+        self.cache[query.lower()] = results
+
+class FullChatDataset(Dataset):
+    def __init__(self, dataset_names=["blended_skill_talk", "conv_ai_2", "social_i_qa"], max_length=256):
+        self.datasets = []
+        
+        for name in dataset_names:
+            try:
+                dataset = load_dataset(name, split="train")
+                self.datasets.append(dataset)
+            except Exception as e:
+                print(f"Failed to load dataset {name}: {e}")
+        
+        self.tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
+        self.tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+        self.max_length = max_length
+        
+    def __len__(self):
+        return sum(len(d) for d in self.datasets)
+    
+    def __getitem__(self, idx):
+        for dataset in self.datasets:
+            if idx < len(dataset):
+                item = dataset[idx]
+                break
+            idx -= len(dataset)
+        
+        if 'dialog' in item:
+            dialog = item['dialog']
+        elif 'messages' in item:
+            dialog = [msg['text'] for msg in item['messages']]
+        else:
+            dialog = [v for k, v in item.items() if isinstance(v, str)]
+        
+        context = " [SEP] ".join(dialog[:-1])
+        response = dialog[-1]
+        
+        inputs = self.tokenizer(
+            context,
+            text_pair=response,
+            max_length=self.max_length,
+            padding='max_length',
+            truncation=True,
+            return_tensors="pt"
+        )
+        
+        return {
+            'input_ids': inputs['input_ids'].flatten(),
+            'attention_mask': inputs['attention_mask'].flatten(),
+            'labels': inputs['input_ids'].flatten()
+        }
+
+class SimpleTransformerModel(nn.Module):
+    def __init__(self, vocab_size, d_model=256, nhead=4, num_layers=3):
+        super().__init__()
+        self.embedding = nn.Embedding(vocab_size, d_model)
+        self.pos_encoder = PositionalEncoding(d_model)
+        encoder_layer = nn.TransformerEncoderLayer(d_model, nhead)
+        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers)
+        self.fc = nn.Linear(d_model, vocab_size)
+        
+    def forward(self, x, mask=None):
+        x = self.embedding(x)
+        x = self.pos_encoder(x)
+        x = self.transformer(x, mask)
+        return self.fc(x)
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model, max_len=500):
+        super().__init__()
+        position = torch.arange(max_len).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
+        pe = torch.zeros(max_len, d_model)
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        self.register_buffer('pe', pe)
+        
+    def forward(self, x):
+        return x + self.pe[:x.size(1)]
+
+class VoiceInterface:
+    def __init__(self):
+        self.recognizer = sr.Recognizer()
+        self.engine = pyttsx3.init()
+        
+    def listen(self) -> Union[str, None]:
+        with sr.Microphone() as source:
+            print("Listening...")
+            audio = self.recognizer.listen(source)
+            try:
+                text = self.recognizer.recognize_google(audio)
+                print(f"You said: {text}")
+                return text
+            except Exception as e:
+                print(f"Error recognizing speech: {e}")
+                return None
+                
+    def speak(self, text: str):
+        print(f"Bot: {text}")
+        self.engine.say(text)
+        self.engine.runAndWait()
+
+class ChatBot:
+    def __init__(self):
+        self.dataset = FullChatDataset()
+        self.model = SimpleTransformerModel(len(self.dataset.tokenizer))
+        self.voice_interface = VoiceInterface()
+        self.web_searcher = WebSearchWrapper()
+        
+    def train(self, epochs=3, lr=3e-4):
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.model = self.model.to(device)
+        criterion = nn.CrossEntropyLoss(ignore_index=0)
+        optimizer = optim.Adam(self.model.parameters(), lr=lr)
+        
+        dataloader = DataLoader(self.dataset, batch_size=8, shuffle=True)
+        
+        for epoch in range(epochs):
+            self.model.train()
+            total_loss = 0
+            pbar = tqdm(dataloader, desc=f"Epoch {epoch+1}/{epochs}")
+            
+            for batch in pbar:
+                inputs = batch['input_ids'].to(device)
+                masks = batch['attention_mask'].to(device)
+                labels = batch['labels'].to(device)
+                
+                optimizer.zero_grad()
+                outputs = self.model(inputs, masks)
+                loss = criterion(outputs.view(-1, outputs.size(-1)), labels.view(-1))
+                loss.backward()
+                optimizer.step()
+                
+                total_loss += loss.item()
+                pbar.set_postfix({'loss': loss.item()})
+            
+            print(f"Epoch {epoch+1} - Avg loss: {total_loss/len(dataloader):.4f}")
+    
+    def generate_response(self, prompt: str, max_length: int = 100, use_web: bool = True) -> str:
+        device = next(self.model.parameters()).device
+        self.model.eval()
+        
+        # Add web context if needed
+        if use_web and self._needs_web_search(prompt):
+            web_results = self.web_searcher.search(prompt)
+            if web_results:
+                prompt = f"Web context: {', '.join(web_results[:3])}. User question: {prompt}"
+        
+        inputs = self.dataset.tokenizer(
+            prompt,
+            return_tensors="pt",
+            max_length=256,
+            truncation=True,
+            padding='max_length'
+        ).to(device)
+        
+        with torch.no_grad():
+            outputs = self.model.generate(
+                input_ids=inputs['input_ids'],
+                attention_mask=inputs['attention_mask'],
+                max_length=max_length,
+                do_sample=True,
+                top_k=50,
+                top_p=0.95,
+                temperature=0.7
+            )
+        
+        response = self.dataset.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return response
+    
+    def _needs_web_search(self, text: str) -> bool:
+        """Determine if a query needs web search"""
+        question_words = ['what', 'when', 'where', 'who', 'why', 'how', 'which', '?']
+        return any(word in text.lower() for word in question_words)