upload

app.py

@@ -1,7 +1,7 @@
 import streamlit as st
 import pandas as pd
 import asyncio
-from llama_models import process_text
+from llama_models import process_text_local
 from dotenv import load_dotenv
 import os
 
@@ -10,6 +10,8 @@ load_dotenv()
 
 # Ensure API key is loaded correctly
 api_key = os.getenv("HUGGINGFACE_API_KEY")
+if api_key is None:
+    raise ValueError("Hugging Face API key is not set. Please add it as a secret in your Hugging Face Space settings.")
 print(f"Hugging Face API Key: {api_key}")
 
 async def process_csv(file):
@@ -28,7 +30,7 @@ async def process_csv(file):
     results = []
     for i, desc in enumerate(descriptions_subset):
         print(f"Processing description {i+1}/{SAMPLE_SIZE}...")
-        result = await process_text(model_name, desc)
+        result = await process_text_local(model_name, desc)
         print(f"Description {i+1} processed. Result: {result[:50]}...")  # Print first 50 characters of the result
         results.append(result)
     

llama_models.py

@@ -1,8 +1,11 @@
 import os
 from transformers import AutoTokenizer, AutoModelForCausalLM  # Ensure correct model class
-import aiohttp
 
 HUGGINGFACE_API_KEY = os.getenv("HUGGINGFACE_API_KEY")
+if HUGGINGFACE_API_KEY is None:
+    raise ValueError("Hugging Face API key is not set. Please add it as a secret in your Hugging Face Space settings.")
+print(f"Using Hugging Face API Key: {HUGGINGFACE_API_KEY}")
+
 model = None
 tokenizer = None
 
@@ -15,21 +18,9 @@ def load_model(model_name):
         print("Model and tokenizer loaded successfully.")
     return tokenizer, model
 
-async def process_text(model_name, text):
+async def process_text_local(model_name, text):
     tokenizer, model = load_model(model_name)
-    prompt = f"Given the following company description, extract key products, geographies, and important keywords:\n\n{text}\n\nProducts, geographies, and keywords:"
-
-    async with aiohttp.ClientSession() as session:
-        print(f"Sending request to model API for text: {text[:50]}...")
-        async with session.post(f"https://api-inference.huggingface.co/models/{model_name}", 
-                                headers={"Authorization": f"Bearer {HUGGINGFACE_API_KEY}"},
-                                json={"inputs": prompt}) as response:
-            print(f"Received response with status code: {response.status}")
-            result = await response.json()
-            print(f"Raw API response: {result}")
-            if isinstance(result, list) and len(result) > 0:
-                return result[0].get('generated_text', '').strip()
-            elif isinstance(result, dict):
-                return result.get('generated_text', '').strip()
-            else:
-                return str(result)
+    inputs = tokenizer.encode(text, return_tensors="pt")
+    outputs = model.generate(inputs, max_length=512)
+    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return result