Upload app.py

app.py

@@ -7,13 +7,32 @@ import gradio as gr
 from parsers import parse_pcap
 from analysis import analyze_calls
 from call_flow import create_call_flow_diagram
-from llm_utils import create_local_pipeline, get_llm_opinion
 
-# 1) Initialize your local Hugging Face pipeline (model)
-#    Adjust model_id to your chosen model on HF.
-MODEL_ID = "google/flan-t5-base"  # "tiiuae/falcon-7b-instruct"
+# We'll adapt llm_utils to load flan-t5-base locally
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
+
+def create_local_pipeline(model_id="google/flan-t5-base"):
+    """
+    Create a local pipeline for Flan-T5, which is a seq2seq model.
+    This should run within ~16GB RAM for 'flan-t5-base'.
+    """
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_id, device_map="auto")
+    # Use text2text-generation for T5-based models
+    return pipeline("text2text-generation", model=model, tokenizer=tokenizer)
+
+# Initialize model pipeline once
+MODEL_ID = "google/flan-t5-base"
 generator = create_local_pipeline(MODEL_ID)
 
+def get_llm_opinion(prompt, generator):
+    """
+    Generate text from the local Flan-T5 pipeline.
+    We use text2text-generation to handle the seq2seq nature of T5.
+    """
+    outputs = generator(prompt, max_length=256, do_sample=True, temperature=0.5)
+    return outputs[0]["generated_text"]
+
 def process_file(pcap_file):
     """
     This function is called when user clicks the 'Analyze File' button.
@@ -64,16 +83,16 @@ def ask_llm_opinion(calls_data, question):
     if not question.strip():
         return "Please enter a question."
 
-    # You might want to create a summary of the calls or pass the entire calls_data.
-    # For brevity, let's just embed a brief mention that calls_data exist:
-    calls_context = "Below is a representation of the calls found in the PCAP:\n"
+    # Summarize the calls for context
+    calls_context = "Below is a brief description of the calls found in the PCAP:\n"
     for call_id, call_obj in calls_data.items():
         calls_context += f"- Call-ID: {call_id}, from_tag: {call_obj.from_tag}, to_tag: {call_obj.to_tag}\n"
 
+    # Build a prompt for T5
     prompt = (
         f"{calls_context}\n"
-        f"User's question: {question}\n"
-        f"Please provide your expert VoIP analysis or advice."
+        f"User's question: {question}\n\n"
+        "Please provide your expert VoIP analysis or advice."
     )
 
     # Query the local pipeline
@@ -87,7 +106,7 @@ def main():
     2) LLM Consultation
     """
     with gr.Blocks() as demo:
-        gr.Markdown("# VoIP Analyzer\nUpload a PCAP/PCAPNG file for SIP/RTP analysis. Then consult an LLM for further insights.")
+        gr.Markdown("# VoIP Analyzer\nUpload a PCAP/PCAPNG file for SIP/RTP analysis. Then consult Flan-T5 for further insights.")
 
         # We keep the calls data in a Gradio State so we can pass it between tabs
         calls_state = gr.State({})