File size: 3,856 Bytes
6f14d8b
 
3a5be9b
6f14d8b
3a5be9b
6f14d8b
3a5be9b
 
 
 
 
 
 
 
 
 
 
 
6f14d8b
3a5be9b
 
6f14d8b
3a5be9b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6f14d8b
3a5be9b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6f14d8b
3a5be9b
 
 
6f14d8b
 
3a5be9b
6f14d8b
 
3a5be9b
 
 
 
 
 
 
 
 
 
6f14d8b
3a5be9b
 
6f14d8b
3a5be9b
 
 
 
 
 
 
 
 
 
 
 
 
 
6f14d8b
3a5be9b
 
 
 
 
6f14d8b
3a5be9b
6f14d8b
 
3a5be9b
 
 
 
 
6f14d8b
 
3a5be9b
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
import logging
from transformers import pipeline
import asyncio

class TextGenerationHandler:
    def __init__(self):
        # Initialize the text generation pipeline
        self.pipe = pipeline(
            "text2text-generation",
            model="google/flan-t5-small",
            max_length=2048,  # Increase max length for longer responses
            num_return_sequences=1
        )
        self.logger = logging.getLogger(__name__)
    
    async def generate_response(self, prompt: str) -> str:
        """
        Generate a complete response using the T5 model pipeline
        
        Args:
            prompt (str): Input text to generate from
            
        Returns:
            str: Generated text output
        """
        try:
            # Break down the generation into sections for better coherence
            sections = [
                "1. Executive Summary",
                "2. Project Scope and Objectives",
                "3. Architecture Overview",
                "4. Component Design",
                "5. Security and Compliance",
                "6. Deployment Strategy",
                "7. Team Requirements",
                "8. Cost Estimates",
                "9. Project Timeline"
            ]
            
            complete_response = []
            for section in sections:
                section_prompt = f"{prompt}\nGenerate content for: {section}"
                self.logger.info(f"Generating section: {section}")
                self.logger.debug(f"Section prompt: {section_prompt}")
                
                output = self.pipe(
                    section_prompt,
                    max_length=512,
                    do_sample=True,
                    temperature=0.7,
                    repetition_penalty=1.2,
                    no_repeat_ngram_size=3
                )
                
                section_text = output[0]['generated_text'].strip()
                self.logger.info(f"Generated text for {section}:\n{section_text}\n")
                
                complete_response.append(f"{section}\n{section_text}")
            
            final_response = "\n\n".join(complete_response)
            self.logger.info(f"Complete response:\n{final_response}")
            return final_response
            
        except Exception as e:
            self.logger.error(f"Error generating text: {str(e)}", exc_info=True)
            raise

    async def stream_response(self, prompt: str):
        """
        Stream the generated response section by section
        
        Args:
            prompt (str): Input text to generate from
            
        Yields:
            dict: Response chunks with type and content
        """
        try:
            # Generate complete response first
            response = await self.generate_response(prompt)
            
            # Stream each section
            accumulated_response = ""
            sections = response.split('\n\n')
            
            for section in sections:
                accumulated_response += section + "\n\n"
                self.logger.debug(f"Streaming section:\n{section}\n")
                
                yield {
                    "type": "content",
                    "content": accumulated_response.strip()
                }
                await asyncio.sleep(0.1)
                    
        except Exception as e:
            self.logger.error(f"Error in stream_response: {str(e)}", exc_info=True)
            yield {
                "type": "error",
                "content": str(e)
            }

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler('f5_model.log')
    ]
)

f5_model = TextGenerationHandler()