ngrigg commited on
Commit
6f15a2e
1 Parent(s): d387575

Securely manage API key and add updated files

Browse files
Files changed (3) hide show
  1. .gitignore +1 -0
  2. llama_models.py +26 -0
  3. requirements.txt +6 -0
.gitignore ADDED
@@ -0,0 +1 @@
 
 
1
+ .env
llama_models.py ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
3
+ import aiohttp
4
+
5
+ HUGGINGFACE_API_KEY = os.getenv("HUGGINGFACE_API_KEY")
6
+
7
+ def load_model(model_name):
8
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
9
+ model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
10
+ return tokenizer, model
11
+
12
+ async def process_text(model_name, text):
13
+ tokenizer, model = load_model(model_name)
14
+ prompt = f"Given the following company description, extract key products, geographies, and important keywords:\n\n{text}\n\nProducts, geographies, and keywords:"
15
+
16
+ async with aiohttp.ClientSession() as session:
17
+ async with session.post(f"https://api-inference.huggingface.co/models/{model_name}",
18
+ headers={"Authorization": f"Bearer {HUGGINGFACE_API_KEY}"},
19
+ json={"inputs": prompt}) as response:
20
+ result = await response.json()
21
+ if isinstance(result, list) and len(result) > 0:
22
+ return result[0].get('generated_text', '').strip()
23
+ elif isinstance(result, dict):
24
+ return result.get('generated_text', '').strip()
25
+ else:
26
+ return str(result)
requirements.txt ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ transformers
2
+ torch
3
+ streamlit
4
+ aiohttp
5
+ pandas
6
+ python-dotenv