Update app.py

app.py

@@ -3,18 +3,21 @@ import torch
 from transformers import AutoTokenizer, T5ForConditionalGeneration, pipeline
 from sentence_transformers import SentenceTransformer, util
 import requests
-import warnings
 import os
-from concurrent.futures import ThreadPoolExecutor
+import warnings
+from transformers import logging
 
-# Set environment variables and suppress warnings
-os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  # Reduce TensorFlow verbosity
-warnings.filterwarnings("ignore", category=FutureWarning)  # Suppress FutureWarnings
-warnings.filterwarnings("ignore", category=UserWarning)    # Suppress UserWarnings
+# Suppress warnings
+warnings.filterwarnings("ignore", category=FutureWarning)
+warnings.filterwarnings("ignore", category=UserWarning)
+warnings.filterwarnings("ignore")
+logging.set_verbosity_error()
 
-GROQ_API_KEY = os.getenv("GROQ_API_KEY")
+# Set API keys and environment variables
+GROQ_API_KEY = os.getenv("GROQ_API_KEY")  # Ensure you set this in Hugging Face Spaces
+os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
 
-# GPT-powered sentence segmentation function
+# Groq API sentence segmentation
 def segment_into_sentences_groq(passage):
     headers = {
         "Authorization": f"Bearer {GROQ_API_KEY}",
@@ -25,146 +28,109 @@ def segment_into_sentences_groq(passage):
         "messages": [
             {
                 "role": "system",
-                "content": "you are to segment the sentence by adding '1!2@3#' at the end of each sentence. Return only the segmented sentences, nothing else."
+                "content": "Segment sentences by adding '1!2@3#' at the end of each sentence."
             },
             {
                 "role": "user",
-                "content": f"Segment this passage into sentences with '1!2@3#' as a delimiter: {passage}"
+                "content": f"Segment the passage: {passage}"
             }
         ],
-        "temperature": 0.7,
-        "max_tokens": 1024
+        "temperature": 1.0,
+        "max_tokens": 8192
     }
-    
     response = requests.post("https://api.groq.com/openai/v1/chat/completions", json=payload, headers=headers)
     if response.status_code == 200:
-        try:
-            segmented_text = response.json()["choices"][0]["message"]["content"]
-            sentences = segmented_text.split("1!2@3#")
-            return [sentence.strip() for sentence in sentences if sentence.strip()]
-        except (KeyError, IndexError):
-            raise ValueError("Unexpected response structure from Groq API.")
+        data = response.json()
+        segmented_text = data.get("choices", [{}])[0].get("message", {}).get("content", "")
+        sentences = segmented_text.split("1!2@3#")
+        return [sentence.strip() for sentence in sentences if sentence.strip()]
     else:
         raise ValueError(f"Groq API error: {response.text}")
 
-
+# Text enhancement class
 class TextEnhancer:
     def __init__(self):
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.executor = ThreadPoolExecutor(max_workers=3)  # Parallel processing pool
-        
-        # Load models
-        self._load_models()
-
-    def _load_models(self):
         self.paraphrase_tokenizer = AutoTokenizer.from_pretrained("prithivida/parrot_paraphraser_on_T5")
         self.paraphrase_model = T5ForConditionalGeneration.from_pretrained("prithivida/parrot_paraphraser_on_T5").to(self.device)
-        
         self.grammar_pipeline = pipeline(
             "text2text-generation",
             model="Grammarly/coedit-large",
             device=0 if self.device == "cuda" else -1
         )
-        
         self.similarity_model = SentenceTransformer('paraphrase-MiniLM-L6-v2').to(self.device)
 
-    def enhance_text(self, text, min_similarity=0.8):
+    def enhance_text(self, text, min_similarity=0.8, max_variations=2):
         sentences = segment_into_sentences_groq(text)
-
-        # Process sentences in parallel
-        results = list(self.executor.map(lambda s: self._process_sentence(s, min_similarity), sentences))
-        
-        # Join enhanced sentences into a single text
-        enhanced_text = ". ".join(results).strip() + "."
-        return enhanced_text
-
-    def _process_sentence(self, sentence, min_similarity):
-        if not sentence.strip():
-            return sentence
-
-        # Generate paraphrases
-        inputs = self.paraphrase_tokenizer(
-            f"paraphrase: {sentence}",
-            return_tensors="pt",
-            padding=True,
-            max_length=150,
-            truncation=True
-        ).to(self.device)
-        
-        outputs = self.paraphrase_model.generate(
-            **inputs,
-            max_length=len(sentence.split()) + 20,
-            num_return_sequences=3,
-            num_beams=3,
-            temperature=0.7
-        )
-        
-        paraphrases = [
-            self.paraphrase_tokenizer.decode(output, skip_special_tokens=True)
-            for output in outputs
-        ]
-
-        # Calculate semantic similarity
-        sentence_embedding = self.similarity_model.encode(sentence, convert_to_tensor=True)
-        paraphrase_embeddings = self.similarity_model.encode(paraphrases, convert_to_tensor=True)
-        similarities = util.cos_sim(sentence_embedding, paraphrase_embeddings).squeeze()
-        
-        # Filter paraphrases by similarity
-        valid_paraphrases = [
-            para for para, sim in zip(paraphrases, similarities)
-            if sim >= min_similarity
-        ]
-        
-        # Grammar correction for the most similar paraphrase
-        if valid_paraphrases:
-            corrected = self.grammar_pipeline(valid_paraphrases[0])[0]["generated_text"]
-            return self._humanize_text(corrected)
-        else:
-            return sentence
-
-    def _humanize_text(self, text):
-        # Introduce minor variations to mimic human-written text
-        import random
-        contractions = {"can't": "cannot", "won't": "will not", "it's": "it is"}
-        words = text.split()
-        text = " ".join([contractions.get(word, word) if random.random() > 0.9 else word for word in words])
-
-        if random.random() > 0.7:
-            text = text.replace(" and ", ", and ")
-        return text
-
-
+        enhanced_sentences = []
+
+        for sentence in sentences:
+            if not sentence.strip():
+                continue
+
+            # Generate paraphrases
+            inputs = self.paraphrase_tokenizer(
+                f"paraphrase: {sentence}",
+                return_tensors="pt",
+                padding=True,
+                max_length=150,
+                truncation=True
+            ).to(self.device)
+            
+            outputs = self.paraphrase_model.generate(
+                **inputs,
+                max_length=150,
+                num_return_sequences=max_variations,
+                num_beams=max_variations
+            )
+            paraphrases = [
+                self.paraphrase_tokenizer.decode(output, skip_special_tokens=True)
+                for output in outputs
+            ]
+
+            # Calculate semantic similarity
+            sentence_embedding = self.similarity_model.encode(sentence)
+            paraphrase_embeddings = self.similarity_model.encode(paraphrases)
+            similarities = util.cos_sim(sentence_embedding, paraphrase_embeddings)
+
+            # Select the most similar paraphrase
+            valid_paraphrases = [
+                para for para, sim in zip(paraphrases, similarities[0])
+                if sim >= min_similarity
+            ]
+            if valid_paraphrases:
+                corrected = self.grammar_pipeline(
+                    valid_paraphrases[0],
+                    max_length=150,
+                    num_return_sequences=1
+                )[0]["generated_text"]
+                enhanced_sentences.append(corrected)
+            else:
+                enhanced_sentences.append(sentence)
+
+        return ". ".join(enhanced_sentences).strip() + "."
+
+# Gradio interface
 def create_interface():
     enhancer = TextEnhancer()
-    
+
     def process_text(text, similarity_threshold):
         try:
             return enhancer.enhance_text(text, min_similarity=similarity_threshold / 100)
         except Exception as e:
             return f"Error: {str(e)}"
-    
-    interface = gr.Interface(
+
+    return gr.Interface(
         fn=process_text,
         inputs=[
-            gr.Textbox(
-                label="Input Text",
-                placeholder="Enter text to enhance...",
-                lines=10
-            ),
-            gr.Slider(
-                minimum=50,
-                maximum=100,
-                value=80,
-                label="Minimum Semantic Similarity (%)"
-            )
+            gr.Textbox(lines=10, placeholder="Enter text to enhance...", label="Input Text"),
+            gr.Slider(50, 100, 80, label="Minimum Semantic Similarity (%)")
         ],
-        outputs=gr.Textbox(label="Enhanced Text", lines=10),
+        outputs=gr.Textbox(lines=10, label="Enhanced Text"),
         title="Text Enhancement System",
-        description="Improve text quality while preserving original meaning.",
+        description="Enhance text quality with semantic preservation."
     )
-    return interface
-
 
 if __name__ == "__main__":
     interface = create_interface()
-    interface.launch()
+    interface.launch(server_name="0.0.0.0", server_port=7860)