Update app.py

app.py

@@ -1,6 +1,6 @@
 """
 DNA-Diffusion Gradio Application
-Interactive DNA sequence generation with slot machine visualization
+Interactive DNA sequence generation with slot machine visualization and protein analysis
 """
 
 import gradio as gr
@@ -9,6 +9,8 @@ import json
 import os
 from typing import Dict, Any, Tuple
 import html
+import requests
+import time
 
 # Configure logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
@@ -45,6 +47,115 @@ if not os.path.exists(HTML_FILE):
 with open(HTML_FILE, "r") as f:
     SLOT_MACHINE_HTML = f.read()
 
+class ProteinAnalyzer:
+    """Handles protein translation and analysis using LLM"""
+    
+    # Genetic code table for DNA to amino acid translation
+    CODON_TABLE = {
+        'TTT': 'F', 'TTC': 'F', 'TTA': 'L', 'TTG': 'L',
+        'TCT': 'S', 'TCC': 'S', 'TCA': 'S', 'TCG': 'S',
+        'TAT': 'Y', 'TAC': 'Y', 'TAA': '*', 'TAG': '*',
+        'TGT': 'C', 'TGC': 'C', 'TGA': '*', 'TGG': 'W',
+        'CTT': 'L', 'CTC': 'L', 'CTA': 'L', 'CTG': 'L',
+        'CCT': 'P', 'CCC': 'P', 'CCA': 'P', 'CCG': 'P',
+        'CAT': 'H', 'CAC': 'H', 'CAA': 'Q', 'CAG': 'Q',
+        'CGT': 'R', 'CGC': 'R', 'CGA': 'R', 'CGG': 'R',
+        'ATT': 'I', 'ATC': 'I', 'ATA': 'I', 'ATG': 'M',
+        'ACT': 'T', 'ACC': 'T', 'ACA': 'T', 'ACG': 'T',
+        'AAT': 'N', 'AAC': 'N', 'AAA': 'K', 'AAG': 'K',
+        'AGT': 'S', 'AGC': 'S', 'AGA': 'R', 'AGG': 'R',
+        'GTT': 'V', 'GTC': 'V', 'GTA': 'V', 'GTG': 'V',
+        'GCT': 'A', 'GCC': 'A', 'GCA': 'A', 'GCG': 'A',
+        'GAT': 'D', 'GAC': 'D', 'GAA': 'E', 'GAG': 'E',
+        'GGT': 'G', 'GGC': 'G', 'GGA': 'G', 'GGG': 'G'
+    }
+    
+    @staticmethod
+    def dna_to_protein(dna_sequence: str) -> str:
+        """Translate DNA sequence to protein sequence"""
+        # Ensure sequence is uppercase
+        dna_sequence = dna_sequence.upper()
+        
+        # Remove any non-DNA characters
+        dna_sequence = ''.join(c for c in dna_sequence if c in 'ATCG')
+        
+        # Translate to protein
+        protein = []
+        for i in range(0, len(dna_sequence) - 2, 3):
+            codon = dna_sequence[i:i+3]
+            if len(codon) == 3:
+                amino_acid = ProteinAnalyzer.CODON_TABLE.get(codon, 'X')
+                if amino_acid == '*':  # Stop codon
+                    break
+                protein.append(amino_acid)
+        
+        return ''.join(protein)
+    
+    @staticmethod
+    def analyze_protein_with_llm(protein_sequence: str, cell_type: str) -> str:
+        """Analyze protein structure and function using Friendli LLM API"""
+        
+        # Get API token from environment
+        token = os.getenv("FRIENDLI_TOKEN")
+        if not token:
+            logger.warning("FRIENDLI_TOKEN not found in environment variables")
+            return "Protein analysis unavailable: API token not configured"
+        
+        try:
+            url = "https://api.friendli.ai/dedicated/v1/chat/completions"
+            headers = {
+                "Authorization": f"Bearer {token}",
+                "Content-Type": "application/json"
+            }
+            
+            # Create prompt for protein analysis
+            prompt = f"""You are a bioinformatics expert. Analyze the following protein sequence and provide insights about its potential structure and function.
+
+Protein sequence: {protein_sequence}
+Cell type context: {cell_type}
+
+Please provide:
+1. Predicted protein family or domain based on sequence patterns
+2. Potential structural features (alpha helices, beta sheets, loops)
+3. Possible biological functions
+4. Relevance to the {cell_type} cell type
+5. Any notable sequence motifs or characteristics
+
+Keep the response concise but informative, suitable for display in a scientific application."""
+            
+            payload = {
+                "model": "dep89a2fld32mcm",
+                "messages": [
+                    {
+                        "role": "system",
+                        "content": "You are a knowledgeable bioinformatics assistant specializing in protein structure and function prediction."
+                    },
+                    {
+                        "role": "user",
+                        "content": prompt
+                    }
+                ],
+                "max_tokens": 1000,
+                "temperature": 0.7,
+                "top_p": 0.8,
+                "stream": False  # Disable streaming for simplicity
+            }
+            
+            response = requests.post(url, json=payload, headers=headers, timeout=30)
+            response.raise_for_status()
+            
+            result = response.json()
+            analysis = result['choices'][0]['message']['content']
+            
+            return analysis
+            
+        except requests.exceptions.RequestException as e:
+            logger.error(f"Failed to analyze protein with LLM: {e}")
+            return f"Protein analysis failed: {str(e)}"
+        except Exception as e:
+            logger.error(f"Unexpected error during protein analysis: {e}")
+            return "Protein analysis unavailable due to an error"
+
 class DNADiffusionApp:
     """Main application class for DNA-Diffusion Gradio interface"""
     
@@ -52,6 +163,7 @@ class DNADiffusionApp:
         self.model = None
         self.model_loading = False
         self.model_error = None
+        self.protein_analyzer = ProteinAnalyzer()
         
     def initialize_model(self):
         """Initialize the DNA-Diffusion model"""
@@ -91,7 +203,6 @@ class DNADiffusionApp:
                 'mock': True
             }
             # Simulate generation time
-            import time
             time.sleep(2.0)
             return sequence, metadata
         
@@ -107,7 +218,28 @@ class DNADiffusionApp:
         """Handle sequence generation request from Gradio"""
         try:
             logger.info(f"Generating sequence for cell type: {cell_type}")
+            
+            # Generate DNA sequence
             sequence, metadata = self.generate_sequence(cell_type, guidance_scale)
+            
+            # Translate to protein
+            logger.info("Translating DNA to protein sequence...")
+            protein_sequence = self.protein_analyzer.dna_to_protein(sequence)
+            
+            # Add protein sequence to metadata
+            metadata['protein_sequence'] = protein_sequence
+            metadata['protein_length'] = len(protein_sequence)
+            
+            # Analyze protein with LLM
+            logger.info("Analyzing protein structure and function...")
+            protein_analysis = self.protein_analyzer.analyze_protein_with_llm(
+                protein_sequence, cell_type
+            )
+            
+            # Add analysis to metadata
+            metadata['protein_analysis'] = protein_analysis
+            
+            logger.info("Generation and analysis complete")
             return sequence, json.dumps(metadata)
         
         except Exception as e: