NEYIRE/modules/docs_assistant.py

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

class DocsAssistant:
    def __init__(self):
        self.model_name = "Arc53/docsgpt-40b-falcon"
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
        self.model = AutoModelForCausalLM.from_pretrained(
            self.model_name,
            torch_dtype=torch.bfloat16,
            device_map="auto"
        )
    
    def search_docs(self, query, doc_file):
        # Process uploaded documentation
        doc_content = self._process_doc_file(doc_file)
        
        # Create prompt with context
        prompt = f"Documentation: {doc_content}\nQuery: {query}"
        
        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
        outputs = self.model.generate(
            **inputs,
            max_length=1024,
            temperature=0.3,
            top_p=0.95
        )
        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    def _process_doc_file(self, file):
        # Implementation for file processing
        pass